diff --git a/templates/components/trustgraph.jsonnet b/templates/components/trustgraph.jsonnet index 833d932b..0c07ea4e 100644 --- a/templates/components/trustgraph.jsonnet +++ b/templates/components/trustgraph.jsonnet @@ -1,6 +1,8 @@ local base = import "base/base.jsonnet"; local images = import "values/images.jsonnet"; local url = import "values/url.jsonnet"; +local minio = import "stores/minio.jsonnet"; +local cassandra = import "stores/cassandra.jsonnet"; { @@ -182,3 +184,6 @@ local url = import "values/url.jsonnet"; } + // Minio and Cassandra are used by the Librarian + + minio + cassandra + diff --git a/templates/generate b/templates/generate index 6f915cda..2569a4c0 100755 --- a/templates/generate +++ b/templates/generate @@ -50,10 +50,13 @@ class Generator: self.templates.joinpath(filename), self.resources.joinpath(dir, filename), self.resources.joinpath(filename), + pathlib.Path(dir).joinpath(filename), ] else: candidates = [ self.templates.joinpath(filename), + pathlib.Path(dir).joinpath(filename), + pathlib.Path(filename), ] try: @@ -86,7 +89,7 @@ class Packager: def __init__(self): self.templates = pathlib.Path("./templates") - self.resources = pathlib.Path("./resources") + self.resources = pathlib.Path("./") def process( self, config, version="0.0.0", platform="docker-compose", diff --git a/templates/stores/milvus.jsonnet b/templates/stores/milvus.jsonnet index cbeb4268..1c3e3734 100644 --- a/templates/stores/milvus.jsonnet +++ b/templates/stores/milvus.jsonnet @@ -1,7 +1,8 @@ local base = import "base/base.jsonnet"; local images = import "values/images.jsonnet"; +local minio = import "stores/minio.jsonnet"; -{ +minio { etcd +: { @@ -47,47 +48,6 @@ local images = import "values/images.jsonnet"; }, - mino +: { - - create:: function(engine) - - local vol = engine.volume("minio-data").with_size("20G"); - - local container = - engine.container("minio") - .with_image(images.minio) - .with_command([ - "minio", - "server", - "/minio_data", - "--console-address", - ":9001", - ]) - .with_environment({ - MINIO_ROOT_USER: "minioadmin", - MINIO_ROOT_PASSWORD: "minioadmin", - }) - .with_limits("0.5", "128M") - .with_reservations("0.25", "128M") - .with_port(9001, 9001, "api") - .with_volume_mount(vol, "/minio_data"); - - local containerSet = engine.containers( - "etcd", [ container ] - ); - - local service = - engine.service(containerSet) - .with_port(9001, 9001, "api"); - - engine.resources([ - vol, - containerSet, - service, - ]) - - }, - milvus +: { create:: function(engine) diff --git a/templates/stores/minio.jsonnet b/templates/stores/minio.jsonnet new file mode 100644 index 00000000..6ef1d96f --- /dev/null +++ b/templates/stores/minio.jsonnet @@ -0,0 +1,49 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; + +{ + + minio +: { + + create:: function(engine) + + local vol = engine.volume("minio-data").with_size("20G"); + + local container = + engine.container("minio") + .with_image(images.minio) + .with_command([ + "minio", + "server", + "/minio_data", + "--console-address", + ":9001", + ]) + .with_environment({ + MINIO_ROOT_USER: "minioadmin", + MINIO_ROOT_PASSWORD: "minioadmin", + }) + .with_limits("0.5", "128M") + .with_reservations("0.25", "128M") + .with_port(9000, 9000, "api") + .with_port(9001, 9001, "console") + .with_volume_mount(vol, "/minio_data"); + + local containerSet = engine.containers( + "etcd", [ container ] + ); + + local service = + engine.service(containerSet) + .with_port(9000, 9000, "api") + .with_port(9001, 9001, "console"); + + engine.resources([ + vol, + containerSet, + service, + ]) + + }, + +} diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index d515a450..3e3d3e75 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -5,7 +5,7 @@ local version = import "version.jsonnet"; pulsar: "docker.io/apachepulsar/pulsar:3.3.1", pulsar_manager: "docker.io/apachepulsar/pulsar-manager:v0.4.0", etcd: "quay.io/coreos/etcd:v3.5.15", - minio: "docker.io/minio/minio:RELEASE.2024-08-17T01-24-54Z", + minio: "docker.io/minio/minio:RELEASE.2025-02-03T21-03-04Z", milvus: "docker.io/milvusdb/milvus:v2.4.9", prometheus: "docker.io/prom/prometheus:v2.53.2", grafana: "docker.io/grafana/grafana:11.1.4", diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index b9dba4fa..3a58d51e 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -3,4 +3,6 @@ from . base_processor import BaseProcessor from . consumer import Consumer from . producer import Producer from . consumer_producer import ConsumerProducer +from . publisher import Publisher +from . subscriber import Subscriber diff --git a/trustgraph-base/trustgraph/base/base_processor.py b/trustgraph-base/trustgraph/base/base_processor.py index 8963e990..d3fd46a9 100644 --- a/trustgraph-base/trustgraph/base/base_processor.py +++ b/trustgraph-base/trustgraph/base/base_processor.py @@ -1,4 +1,5 @@ +import asyncio import os import argparse import pulsar @@ -30,6 +31,7 @@ class BaseProcessor: pulsar_host = params.get("pulsar_host", self.default_pulsar_host) pulsar_api_key = params.get("pulsar_api_key", None) + pulsar_listener = params.get("pulsar_listener", None) log_level = params.get("log_level", LogLevel.INFO) self.pulsar_host = pulsar_host @@ -43,9 +45,12 @@ class BaseProcessor: else: self.client = pulsar.Client( pulsar_host, + listener_name=pulsar_listener, logger=pulsar.ConsoleLogger(log_level.to_pulsar()) ) + self.pulsar_listener = pulsar_listener + def __del__(self): if hasattr(self, "client"): @@ -67,6 +72,11 @@ class BaseProcessor: help=f'Pulsar API key', ) + parser.add_argument( + '--pulsar-listener', + help=f'Pulsar listener (default: none)', + ) + parser.add_argument( '-l', '--log-level', type=LogLevel, @@ -89,11 +99,20 @@ class BaseProcessor: help=f'Pulsar host (default: 8000)', ) - def run(self): + async def start(self): + pass + + async def run(self): raise RuntimeError("Something should have implemented the run method") @classmethod - def start(cls, prog, doc): + async def launch_async(cls, args): + p = cls(**args) + await p.start() + await p.run() + + @classmethod + def launch(cls, prog, doc): parser = argparse.ArgumentParser( prog=prog, @@ -114,8 +133,7 @@ class BaseProcessor: try: - p = cls(**args) - p.run() + asyncio.run(cls.launch_async(args)) except KeyboardInterrupt: print("Keyboard interrupt.") @@ -133,3 +151,4 @@ class BaseProcessor: print("Will retry...", flush=True) time.sleep(4) + diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index 521dd3c1..175f1fd7 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -1,4 +1,5 @@ +import asyncio from pulsar.schema import JsonSchema import pulsar from prometheus_client import Histogram, Info, Counter, Enum @@ -75,7 +76,7 @@ class Consumer(BaseProcessor): print("Initialised consumer.", flush=True) - def run(self): + async def run(self): __class__.state_metric.state('running') @@ -104,7 +105,7 @@ class Consumer(BaseProcessor): try: with __class__.request_metric.time(): - self.handle(msg) + await self.handle(msg) # Acknowledge successful processing of the message self.consumer.acknowledge(msg) diff --git a/trustgraph-base/trustgraph/base/consumer_producer.py b/trustgraph-base/trustgraph/base/consumer_producer.py index be9915ce..1006f9b5 100644 --- a/trustgraph-base/trustgraph/base/consumer_producer.py +++ b/trustgraph-base/trustgraph/base/consumer_producer.py @@ -42,7 +42,7 @@ class ConsumerProducer(Consumer): print("Initialised consumer/producer.") - def send(self, msg, properties={}): + async def send(self, msg, properties={}): self.producer.send(msg, properties) __class__.output_metric.inc() diff --git a/trustgraph-base/trustgraph/base/producer.py b/trustgraph-base/trustgraph/base/producer.py index 84d7fc99..bc2d7791 100644 --- a/trustgraph-base/trustgraph/base/producer.py +++ b/trustgraph-base/trustgraph/base/producer.py @@ -37,7 +37,7 @@ class Producer(BaseProcessor): chunking_enabled=True, ) - def send(self, msg, properties={}): + async def send(self, msg, properties={}): self.producer.send(msg, properties) __class__.output_metric.inc() diff --git a/trustgraph-flow/trustgraph/gateway/publisher.py b/trustgraph-base/trustgraph/base/publisher.py similarity index 82% rename from trustgraph-flow/trustgraph/gateway/publisher.py rename to trustgraph-base/trustgraph/base/publisher.py index 5dfb882b..2b50ed78 100644 --- a/trustgraph-flow/trustgraph/gateway/publisher.py +++ b/trustgraph-base/trustgraph/base/publisher.py @@ -15,14 +15,22 @@ class Publisher: self.q = queue.Queue(maxsize=max_size) self.chunking_enabled = chunking_enabled self.listener_name = listener + self.running = True def start(self): self.task = threading.Thread(target=self.run) self.task.start() + def stop(self): + self.running = False + + def join(self): + self.stop() + self.task.join() + def run(self): - while True: + while self.running: try: @@ -44,9 +52,12 @@ class Publisher: chunking_enabled=self.chunking_enabled, ) - while True: + while self.running: - id, item = self.q.get() + try: + id, item = self.q.get(timeout=0.5) + except queue.Empty: + continue if id: producer.send(item, { "id": id }) @@ -61,3 +72,5 @@ class Publisher: def send(self, id, msg): self.q.put((id, msg)) + + diff --git a/trustgraph-flow/trustgraph/gateway/subscriber.py b/trustgraph-base/trustgraph/base/subscriber.py similarity index 90% rename from trustgraph-flow/trustgraph/gateway/subscriber.py rename to trustgraph-base/trustgraph/base/subscriber.py index 60406f6b..33ecea3a 100644 --- a/trustgraph-flow/trustgraph/gateway/subscriber.py +++ b/trustgraph-base/trustgraph/base/subscriber.py @@ -19,14 +19,21 @@ class Subscriber: self.max_size = max_size self.lock = threading.Lock() self.listener_name = listener + self.running = True def start(self): self.task = threading.Thread(target=self.run) self.task.start() + def stop(self): + self.running = False + + def join(self): + self.task.join() + def run(self): - while True: + while self.running: try: @@ -50,7 +57,7 @@ class Subscriber: schema=self.schema, ) - while True: + while self.running: msg = consumer.receive() @@ -68,12 +75,14 @@ class Subscriber: if id in self.q: try: + # FIXME: Timeout means data goes missing self.q[id].put(value, timeout=0.5) except: pass for q in self.full.values(): try: + # FIXME: Timeout means data goes missing q.put(value, timeout=0.5) except: pass diff --git a/trustgraph-base/trustgraph/exceptions.py b/trustgraph-base/trustgraph/exceptions.py index afe72ccc..09f098df 100644 --- a/trustgraph-base/trustgraph/exceptions.py +++ b/trustgraph-base/trustgraph/exceptions.py @@ -8,3 +8,6 @@ class LlmError(Exception): class ParseError(Exception): pass +class RequestError(Exception): + pass + diff --git a/trustgraph-base/trustgraph/schema/__init__.py b/trustgraph-base/trustgraph/schema/__init__.py index be41b670..9c44a743 100644 --- a/trustgraph-base/trustgraph/schema/__init__.py +++ b/trustgraph-base/trustgraph/schema/__init__.py @@ -10,5 +10,6 @@ from . retrieval import * from . metadata import * from . agent import * from . lookup import * +from . library import * diff --git a/trustgraph-base/trustgraph/schema/library.py b/trustgraph-base/trustgraph/schema/library.py new file mode 100644 index 00000000..8ab88842 --- /dev/null +++ b/trustgraph-base/trustgraph/schema/library.py @@ -0,0 +1,56 @@ + +from pulsar.schema import Record, Bytes, String, Array +from . types import Triple +from . topic import topic +from . types import Error +from . metadata import Metadata +from . documents import Document, TextDocument + +# 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 DocumentPackage(Record): + metadata = Array(Triple()) + document = Bytes() + kind = String() + user = String() + collection = String() + +class DocumentInfo(Record): + metadata = Array(Triple()) + kind = String() + user = String() + collection = String() + +class Criteria(Record): + key = String() + value = String() + operator = String() + +class LibrarianRequest(Record): + operation = String() + id = String() + document = DocumentPackage() + user = String() + collection = String() + criteria = Array(Criteria()) + +class LibrarianResponse(Record): + error = Error() + document = DocumentPackage() + info = Array(DocumentInfo()) + +librarian_request_queue = topic( + 'librarian', kind='non-persistent', namespace='request' +) +librarian_response_queue = topic( + 'librarian', kind='non-persistent', namespace='response', +) + diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index f46e35fb..3bca324d 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -109,7 +109,7 @@ class Processor(ConsumerProducer): print("Initialised", flush=True) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -341,5 +341,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-embeddings-hf/trustgraph/embeddings/hf/hf.py b/trustgraph-embeddings-hf/trustgraph/embeddings/hf/hf.py index 4b3b39c1..5470333b 100755 --- a/trustgraph-embeddings-hf/trustgraph/embeddings/hf/hf.py +++ b/trustgraph-embeddings-hf/trustgraph/embeddings/hf/hf.py @@ -40,7 +40,7 @@ class Processor(ConsumerProducer): self.embeddings = HuggingFaceEmbeddings(model_name=model) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -96,5 +96,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/scripts/librarian b/trustgraph-flow/scripts/librarian new file mode 100755 index 00000000..9f6458ab --- /dev/null +++ b/trustgraph-flow/scripts/librarian @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.librarian import run + +run() + diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index e56aa694..fe167e90 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -49,6 +49,7 @@ setuptools.setup( "langchain-community", "langchain-core", "langchain-text-splitters", + "minio", "neo4j", "ollama", "openai", @@ -78,8 +79,8 @@ setuptools.setup( "scripts/de-write-qdrant", "scripts/document-embeddings", "scripts/document-rag", - "scripts/embeddings-ollama", "scripts/embeddings-fastembed", + "scripts/embeddings-ollama", "scripts/ge-query-milvus", "scripts/ge-query-pinecone", "scripts/ge-query-qdrant", @@ -91,6 +92,7 @@ setuptools.setup( "scripts/kg-extract-definitions", "scripts/kg-extract-relationships", "scripts/kg-extract-topics", + "scripts/librarian", "scripts/metering", "scripts/object-extract-row", "scripts/oe-write-milvus", diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index ccfa1cf8..7544f008 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -191,7 +191,7 @@ class Processor(ConsumerProducer): return json.loads(json_str) - def handle(self, msg): + async def handle(self, msg): try: @@ -231,7 +231,7 @@ class Processor(ConsumerProducer): observation=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) def observe(x): @@ -244,7 +244,7 @@ class Processor(ConsumerProducer): observation=x, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) act = self.agent.react(v.question, history, think, observe) @@ -260,7 +260,7 @@ class Processor(ConsumerProducer): thought=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -283,7 +283,7 @@ class Processor(ConsumerProducer): ] ) - self.recursive_input.send(r, properties={"id": id}) + await self.recursive_input.send(r, properties={"id": id}) print("Done.", flush=True) @@ -303,7 +303,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) @staticmethod def add_args(parser): @@ -379,5 +379,5 @@ description.''' def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py index 694ced70..82f333b5 100755 --- a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py @@ -52,7 +52,7 @@ class Processor(ConsumerProducer): is_separator_regex=False, ) - def handle(self, msg): + async def handle(self, msg): v = msg.value() print(f"Chunking {v.metadata.id}...", flush=True) @@ -70,7 +70,7 @@ class Processor(ConsumerProducer): __class__.chunk_metric.observe(len(chunk.page_content)) - self.send(r) + await self.send(r) print("Done.", flush=True) @@ -98,5 +98,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/chunking/token/chunker.py b/trustgraph-flow/trustgraph/chunking/token/chunker.py index dccd9c89..c625b48c 100755 --- a/trustgraph-flow/trustgraph/chunking/token/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/token/chunker.py @@ -51,7 +51,7 @@ class Processor(ConsumerProducer): chunk_overlap=chunk_overlap, ) - def handle(self, msg): + async def handle(self, msg): v = msg.value() print(f"Chunking {v.metadata.id}...", flush=True) @@ -69,7 +69,7 @@ class Processor(ConsumerProducer): __class__.chunk_metric.observe(len(chunk.page_content)) - self.send(r) + await self.send(r) print("Done.", flush=True) @@ -97,5 +97,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index 38ac9257..5e5e3612 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -39,7 +39,7 @@ class Processor(ConsumerProducer): print("PDF inited") - def handle(self, msg): + async def handle(self, msg): print("PDF message received") @@ -64,7 +64,7 @@ class Processor(ConsumerProducer): text=page.page_content.encode("utf-8"), ) - self.send(r) + await self.send(r) print("Done.", flush=True) @@ -78,5 +78,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py index 5fe8a23e..70f53e07 100755 --- a/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py +++ b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py @@ -53,7 +53,7 @@ class Processor(ConsumerProducer): subscriber=module + "-emb", ) - def handle(self, msg): + async def handle(self, msg): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) @@ -74,7 +74,7 @@ class Processor(ConsumerProducer): chunks=embeds, ) - self.producer.send(r) + await self.send(r) except Exception as e: print("Exception:", e, flush=True) @@ -106,5 +106,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py index 635387b8..253d77ac 100755 --- a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py @@ -41,7 +41,7 @@ class Processor(ConsumerProducer): self.embeddings = TextEmbedding(model_name = model) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -65,7 +65,7 @@ class Processor(ConsumerProducer): error=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -85,5 +85,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py index e4d1646e..2cbe9907 100755 --- a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py +++ b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py @@ -52,7 +52,7 @@ class Processor(ConsumerProducer): subscriber=module + "-emb", ) - def handle(self, msg): + async def handle(self, msg): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) @@ -77,7 +77,7 @@ class Processor(ConsumerProducer): entities=entities, ) - self.producer.send(r) + await self.send(r) except Exception as e: print("Exception:", e, flush=True) @@ -109,5 +109,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py index 5baf64aa..f9dd87a9 100755 --- a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py @@ -45,7 +45,7 @@ class Processor(ConsumerProducer): self.client = Client(host=ollama) self.model = model - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -67,7 +67,7 @@ class Processor(ConsumerProducer): error=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -93,5 +93,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/external/wikipedia/service.py b/trustgraph-flow/trustgraph/external/wikipedia/service.py index 932e1213..a269c8ce 100644 --- a/trustgraph-flow/trustgraph/external/wikipedia/service.py +++ b/trustgraph-flow/trustgraph/external/wikipedia/service.py @@ -39,7 +39,7 @@ class Processor(ConsumerProducer): self.url = url - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -60,7 +60,7 @@ class Processor(ConsumerProducer): text=resp ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -75,7 +75,7 @@ class Processor(ConsumerProducer): ), text=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -98,5 +98,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index fe824378..47c99802 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py @@ -97,15 +97,15 @@ class Processor(ConsumerProducer): return self.prompt.request_definitions(chunk) - def emit_edges(self, metadata, triples): + async def emit_edges(self, metadata, triples): t = Triples( metadata=metadata, triples=triples, ) - self.producer.send(t) + await self.send(t) - def emit_ecs(self, metadata, entities): + async def emit_ecs(self, metadata, entities): t = EntityContexts( metadata=metadata, @@ -113,7 +113,7 @@ class Processor(ConsumerProducer): ) self.ec_prod.send(t) - def handle(self, msg): + async def handle(self, msg): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) @@ -172,7 +172,7 @@ class Processor(ConsumerProducer): entities.append(ec) - self.emit_edges( + await self.emit_edges( Metadata( id=v.metadata.id, metadata=[], @@ -182,7 +182,7 @@ class Processor(ConsumerProducer): triples ) - self.emit_ecs( + await self.emit_ecs( Metadata( id=v.metadata.id, metadata=[], @@ -225,5 +225,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py index fae6e07b..2f293527 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py @@ -72,15 +72,15 @@ class Processor(ConsumerProducer): return self.prompt.request_relationships(chunk) - def emit_edges(self, metadata, triples): + async def emit_edges(self, metadata, triples): t = Triples( metadata=metadata, triples=triples, ) - self.producer.send(t) + await self.send(t) - def handle(self, msg): + async def handle(self, msg): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) @@ -167,7 +167,7 @@ class Processor(ConsumerProducer): o=Value(value=v.metadata.id, is_uri=True) )) - self.emit_edges( + await self.emit_edges( Metadata( id=v.metadata.id, metadata=[], @@ -204,5 +204,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py index b68bb3ea..7424abe2 100755 --- a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py @@ -70,15 +70,15 @@ class Processor(ConsumerProducer): return self.prompt.request_topics(chunk) - def emit_edge(self, metadata, s, p, o): + async def emit_edge(self, metadata, s, p, o): t = Triples( metadata=metadata, triples=[Triple(s=s, p=p, o=o)], ) - self.producer.send(t) + await self.send(t) - def handle(self, msg): + async def handle(self, msg): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) @@ -105,7 +105,9 @@ class Processor(ConsumerProducer): s_value = Value(value=str(s_uri), is_uri=True) o_value = Value(value=str(o), is_uri=False) - self.emit_edge(v. metadata, s_value, DEFINITION_VALUE, o_value) + await self.emit_edge( + v.metadata, s_value, DEFINITION_VALUE, o_value + ) except Exception as e: print("Exception: ", e, flush=True) @@ -134,5 +136,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/extract/object/row/extract.py b/trustgraph-flow/trustgraph/extract/object/row/extract.py index 0e2a5701..db02a651 100755 --- a/trustgraph-flow/trustgraph/extract/object/row/extract.py +++ b/trustgraph-flow/trustgraph/extract/object/row/extract.py @@ -130,16 +130,16 @@ class Processor(ConsumerProducer): t = Rows( metadata=metadata, row_schema=self.row_schema, rows=rows ) - self.producer.send(t) + await self.producer.send(t) def emit_vec(self, metadata, name, vec, key_name, key): r = ObjectEmbeddings( metadata=metadata, vectors=vec, name=name, key_name=key_name, id=key ) - self.vec_prod.send(r) + await self.vec_prod.send(r) - def handle(self, msg): + async def handle(self, msg): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) @@ -217,5 +217,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py index 6de41a35..7d1e6496 100644 --- a/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py @@ -7,8 +7,8 @@ from aiohttp import WSMsgType from .. schema import Metadata from .. schema import DocumentEmbeddings, ChunkEmbeddings from .. schema import document_embeddings_store_queue +from .. base import Publisher -from . publisher import Publisher from . socket import SocketEndpoint from . serialize import to_subgraph @@ -61,6 +61,6 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint): ], ) - self.publisher.send(None, elt) + await self.publisher.send(None, elt) running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py index 252a4926..efa8ec84 100644 --- a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py @@ -6,8 +6,8 @@ import uuid from .. schema import DocumentEmbeddings from .. schema import document_embeddings_store_queue +from .. base import Subscriber -from . subscriber import Subscriber from . socket import SocketEndpoint from . serialize import serialize_document_embeddings diff --git a/trustgraph-flow/trustgraph/gateway/endpoint.py b/trustgraph-flow/trustgraph/gateway/endpoint.py index 1f38c489..5005463c 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint.py @@ -5,8 +5,8 @@ from aiohttp import web import uuid import logging -from . publisher import Publisher -from . subscriber import Subscriber +from .. base import Publisher +from .. base import Subscriber logger = logging.getLogger("endpoint") logger.setLevel(logging.INFO) diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py index b2d5b80c..623c0c3e 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py @@ -7,8 +7,8 @@ from aiohttp import WSMsgType from .. schema import Metadata from .. schema import GraphEmbeddings, EntityEmbeddings from .. schema import graph_embeddings_store_queue +from .. base import Publisher -from . publisher import Publisher from . socket import SocketEndpoint from . serialize import to_subgraph, to_value @@ -61,7 +61,7 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint): ] ) - self.publisher.send(None, elt) + await self.publisher.send(None, elt) running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py index c2613144..ddbe0d38 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py @@ -6,8 +6,8 @@ import uuid from .. schema import GraphEmbeddings from .. schema import graph_embeddings_store_queue +from .. base import Subscriber -from . subscriber import Subscriber from . socket import SocketEndpoint from . serialize import serialize_graph_embeddings diff --git a/trustgraph-flow/trustgraph/gateway/librarian.py b/trustgraph-flow/trustgraph/gateway/librarian.py new file mode 100644 index 00000000..78de7970 --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/librarian.py @@ -0,0 +1,57 @@ + +from .. schema import LibrarianRequest, LibrarianResponse, Triples +from .. schema import librarian_request_queue +from .. schema import librarian_response_queue + +from . endpoint import ServiceEndpoint +from . requestor import ServiceRequestor +from . serialize import serialize_document_package, serialize_document_info +from . serialize import to_document_package, to_document_info, to_criteria + +class LibrarianRequestor(ServiceRequestor): + def __init__(self, pulsar_host, timeout, auth): + + super(LibrarianRequestor, self).__init__( + pulsar_host=pulsar_host, + request_queue=librarian_request_queue, + response_queue=librarian_response_queue, + request_schema=LibrarianRequest, + response_schema=LibrarianResponse, + timeout=timeout, + ) + + def to_request(self, body): + + if "document" in body: + dp = to_document_package(body["document"]) + else: + dp = None + + if "criteria" in body: + criteria = to_criteria(body["criteria"]) + else: + criteria = None + + limit = int(body.get("limit", 10000)) + + return LibrarianRequest( + operation = body.get("operation", None), + id = body.get("id", None), + document = dp, + user = body.get("user", None), + collection = body.get("collection", None), + criteria = criteria, + ) + + def from_response(self, message): + + response = {} + + if message.document: + response["document"] = serialize_document_package(message.document) + + if message.info: + response["info"] = serialize_document_info(message.info) + + return response, True + diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py index 9216f681..fc733e16 100644 --- a/trustgraph-flow/trustgraph/gateway/requestor.py +++ b/trustgraph-flow/trustgraph/gateway/requestor.py @@ -4,8 +4,8 @@ from pulsar.schema import JsonSchema import uuid import logging -from . publisher import Publisher -from . subscriber import Subscriber +from .. base import Publisher +from .. base import Subscriber logger = logging.getLogger("requestor") logger.setLevel(logging.INFO) @@ -71,7 +71,10 @@ class ServiceRequestor: raise RuntimeError("Timeout") if resp.error: - err = { "error": resp.error.message } + err = { "error": { + "type": resp.error.type, + "message": resp.error.message, + } } if responder: await responder(err, True) return err @@ -90,7 +93,10 @@ class ServiceRequestor: logging.error(f"Exception: {e}") - err = { "error": str(e) } + err = { "error": { + "type": "gateway-error", + "message": str(e), + } } if responder: await responder(err, True) return err diff --git a/trustgraph-flow/trustgraph/gateway/sender.py b/trustgraph-flow/trustgraph/gateway/sender.py index f232780f..1a66b7e1 100644 --- a/trustgraph-flow/trustgraph/gateway/sender.py +++ b/trustgraph-flow/trustgraph/gateway/sender.py @@ -6,7 +6,7 @@ from pulsar.schema import JsonSchema import uuid import logging -from . publisher import Publisher +from .. base import Publisher logger = logging.getLogger("sender") logger.setLevel(logging.INFO) diff --git a/trustgraph-flow/trustgraph/gateway/serialize.py b/trustgraph-flow/trustgraph/gateway/serialize.py index 40b6efc5..552105c3 100644 --- a/trustgraph-flow/trustgraph/gateway/serialize.py +++ b/trustgraph-flow/trustgraph/gateway/serialize.py @@ -1,4 +1,7 @@ -from .. schema import Value, Triple + +import base64 + +from .. schema import Value, Triple, DocumentPackage, DocumentInfo def to_value(x): return Value(value=x["v"], is_uri=x["e"]) @@ -77,3 +80,69 @@ def serialize_document_embeddings(message): ], } +def serialize_document_package(message): + + ret = {} + + if message.metadata: + ret["metadata"] = serialize_subgraph(message.metdata) + + if message.document: + blob = base64.b64encode( + message.document.encode("utf-8") + ).decode("utf-8") + ret["document"] = blob + + if message.kind: + ret["kind"] = message.kind + + if message.user: + ret["user"] = message.user + + if message.collection: + ret["collection"] = message.collection + + return ret + +def serialize_document_info(message): + + ret = {} + + if message.metadata: + ret["metadata"] = serialize_subgraph(message.metdata) + + if message.kind: + ret["kind"] = message.kind + + if message.user: + ret["user"] = message.user + + if message.collection: + ret["collection"] = message.collection + + return ret + +def to_document_package(x): + + return DocumentPackage( + metadata = to_subgraph(x["metadata"]), + document = base64.b64decode(x["document"].encode("utf-8")), + kind = x.get("kind", None), + user = x.get("user", None), + collection = x.get("collection", None), + ) + +def to_document_info(x): + + return DocumentInfo( + metadata = to_subgraph(x["metadata"]), + kind = x.get("kind", None), + user = x.get("user", None), + collection = x.get("collection", None), + ) + +def to_criteria(x): + return [ + Critera(v["key"], v["value"], v["operator"]) + for v in x + ] diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index 5f131c2b..711d85fd 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -26,8 +26,6 @@ from .. log_level import LogLevel from . serialize import to_subgraph from . running import Running -from . publisher import Publisher -from . subscriber import Subscriber from . text_completion import TextCompletionRequestor from . prompt import PromptRequestor from . graph_rag import GraphRagRequestor @@ -39,6 +37,7 @@ from . encyclopedia import EncyclopediaRequestor from . agent import AgentRequestor from . dbpedia import DbpediaRequestor from . internet_search import InternetSearchRequestor +from . librarian import LibrarianRequestor from . triples_stream import TriplesStreamEndpoint from . graph_embeddings_stream import GraphEmbeddingsStreamEndpoint from . document_embeddings_stream import DocumentEmbeddingsStreamEndpoint @@ -125,6 +124,10 @@ class Api: pulsar_host=self.pulsar_host, timeout=self.timeout, auth = self.auth, pulsar_api_key=self.pulsar_api_key, ), + "librarian": LibrarianRequestor( + pulsar_host=self.pulsar_host, timeout=self.timeout, + auth = self.auth, + ), "encyclopedia": EncyclopediaRequestor( pulsar_host=self.pulsar_host, timeout=self.timeout, auth = self.auth, pulsar_api_key=self.pulsar_api_key, @@ -179,6 +182,10 @@ class Api: endpoint_path = "/api/v1/agent", auth=self.auth, requestor = self.services["agent"], ), + ServiceEndpoint( + endpoint_path = "/api/v1/librarian", auth=self.auth, + requestor = self.services["librarian"], + ), ServiceEndpoint( endpoint_path = "/api/v1/encyclopedia", auth=self.auth, requestor = self.services["encyclopedia"], diff --git a/trustgraph-flow/trustgraph/gateway/triples_load.py b/trustgraph-flow/trustgraph/gateway/triples_load.py index 05d329d0..ccee3009 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_load.py +++ b/trustgraph-flow/trustgraph/gateway/triples_load.py @@ -7,8 +7,8 @@ from aiohttp import WSMsgType from .. schema import Metadata from .. schema import Triples from .. schema import triples_store_queue +from .. base import Publisher -from . publisher import Publisher from . socket import SocketEndpoint from . serialize import to_subgraph @@ -53,7 +53,7 @@ class TriplesLoadEndpoint(SocketEndpoint): triples=to_subgraph(data["triples"]), ) - self.publisher.send(None, elt) + await self.publisher.send(None, elt) running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/triples_stream.py b/trustgraph-flow/trustgraph/gateway/triples_stream.py index b98ba437..4202620a 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_stream.py +++ b/trustgraph-flow/trustgraph/gateway/triples_stream.py @@ -6,8 +6,8 @@ import uuid from .. schema import Triples from .. schema import triples_store_queue +from .. base import Subscriber -from . subscriber import Subscriber from . socket import SocketEndpoint from . serialize import serialize_triples diff --git a/trustgraph-flow/trustgraph/librarian/__init__.py b/trustgraph-flow/trustgraph/librarian/__init__.py new file mode 100644 index 00000000..ba844705 --- /dev/null +++ b/trustgraph-flow/trustgraph/librarian/__init__.py @@ -0,0 +1,3 @@ + +from . service import * + diff --git a/trustgraph-flow/trustgraph/librarian/__main__.py b/trustgraph-flow/trustgraph/librarian/__main__.py new file mode 100755 index 00000000..e9136855 --- /dev/null +++ b/trustgraph-flow/trustgraph/librarian/__main__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from . service import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph-flow/trustgraph/librarian/blob_store.py b/trustgraph-flow/trustgraph/librarian/blob_store.py new file mode 100644 index 00000000..5cffef18 --- /dev/null +++ b/trustgraph-flow/trustgraph/librarian/blob_store.py @@ -0,0 +1,51 @@ +from .. schema import LibrarianRequest, LibrarianResponse, Error +from .. knowledge import hash +from .. exceptions import RequestError + +from minio import Minio +import time +import io + +class BlobStore: + + def __init__( + self, + minio_host, minio_access_key, minio_secret_key, bucket_name, + ): + + + self.minio = Minio( + minio_host, + access_key = minio_access_key, + secret_key = minio_secret_key, + secure = False, + ) + + self.bucket_name = bucket_name + + print("Connected to minio", flush=True) + + self.ensure_bucket() + + def ensure_bucket(self): + + # Make the bucket if it doesn't exist. + found = self.minio.bucket_exists(self.bucket_name) + if not found: + self.minio.make_bucket(self.bucket_name) + print("Created bucket", self.bucket_name, flush=True) + else: + print("Bucket", self.bucket_name, "already exists", flush=True) + + def add(self, object_id, blob, kind): + + # FIXME: Loop retry + self.minio.put_object( + bucket_name = self.bucket_name, + object_name = "doc/" + str(object_id), + length = len(blob), + data = io.BytesIO(blob), + content_type = kind, + ) + + print("Add blob complete", flush=True) diff --git a/trustgraph-flow/trustgraph/librarian/librarian.py b/trustgraph-flow/trustgraph/librarian/librarian.py new file mode 100644 index 00000000..08010f3f --- /dev/null +++ b/trustgraph-flow/trustgraph/librarian/librarian.py @@ -0,0 +1,56 @@ +from .. schema import LibrarianRequest, LibrarianResponse, Error +from .. knowledge import hash +from .. exceptions import RequestError +from . table_store import TableStore +from . blob_store import BlobStore + +import uuid + +class Librarian: + + def __init__( + self, + cassandra_host, cassandra_user, cassandra_password, + minio_host, minio_access_key, minio_secret_key, + bucket_name, keyspace, load_document, load_text, + ): + + self.blob_store = BlobStore( + minio_host, minio_access_key, minio_secret_key, bucket_name + ) + + self.table_store = TableStore( + cassandra_host, cassandra_user, cassandra_password, keyspace + ) + + self.load_document = load_document + self.load_text = load_text + + async def add(self, id, document): + + if document.kind not in ( + "text/plain", "application/pdf" + ): + raise RequestError("Invalid document kind: " + document.kind) + + # Create object ID as a hash of the document + object_id = uuid.UUID(hash(document.document)) + + self.blob_store.add(object_id, document.document, document.kind) + + self.table_store.add(object_id, document) + + if document.kind == "application/pdf": + await self.load_document(id, document) + elif document.kind == "text/plain": + await self.load_text(id, document) + + print("Add complete", flush=True) + + return LibrarianResponse( + error = None, + document = None, + info = None, + ) + + diff --git a/trustgraph-flow/trustgraph/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py new file mode 100755 index 00000000..ade4ca38 --- /dev/null +++ b/trustgraph-flow/trustgraph/librarian/service.py @@ -0,0 +1,354 @@ + +""" +Librarian service, manages documents in collections +""" + +from functools import partial +import asyncio +import threading +import queue + +from pulsar.schema import JsonSchema + +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 = ".".join(__name__.split(".")[1:-1]) + +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" +default_cassandra_host = "cassandra" + +bucket_name = "library" + +# FIXME: How to ensure this doesn't conflict with other usage? +keyspace = "librarian" + +class Processor(ConsumerProducer): + + def __init__(self, **params): + + self.running = True + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + + minio_host = params.get("minio_host", default_minio_host) + minio_access_key = params.get( + "minio_access_key", + default_minio_access_key + ) + minio_secret_key = params.get( + "minio_secret_key", + default_minio_secret_key + ) + + cassandra_host = params.get("cassandra_host", default_cassandra_host) + 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, + "minio_host": minio_host, + "minio_access_key": minio_access_key, + "cassandra_host": cassandra_host, + "cassandra_user": cassandra_user, + } + ) + + self.document_load = Publisher( + self.pulsar_host, document_load_queue, JsonSchema(Document), + listener=self.pulsar_listener, + ) + + self.text_load = Publisher( + self.pulsar_host, text_load_queue, JsonSchema(TextDocument), + listener=self.pulsar_listener, + ) + + self.triples_load = Subscriber( + self.pulsar_host, triples_store_queue, + "librarian", "librarian", + schema=JsonSchema(Triples), + listener=self.pulsar_listener, + ) + + self.triples_reader = threading.Thread(target=self.receive_triples) + + self.librarian = Librarian( + cassandra_host = cassandra_host.split(","), + cassandra_user = cassandra_user, + cassandra_password = cassandra_password, + minio_host = minio_host, + minio_access_key = minio_access_key, + minio_secret_key = minio_secret_key, + bucket_name = bucket_name, + keyspace = keyspace, + load_document = self.load_document, + load_text = self.load_text, + ) + + print("Initialised.", flush=True) + + async def start(self): + + self.document_load.start() + self.text_load.start() + self.triples_load.start() + + self.triples_sub = self.triples_load.subscribe_all("x") + + self.triples_reader.start() + + def receive_triples(self): + + print("Receive triples!") + + while self.running: + try: + msg = self.triples_sub.get(timeout=1) + except queue.Empty: + print("Tick") + continue + + print(msg) + + print("BYE") + + def __del__(self): + + self.running = False + + if hasattr(self, "triples_sub"): + self.triples_sub.unsubscribe_all("x") + + 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_load"): + self.triples_load.stop() + self.triples_load.join() + + async def load_document(self, id, document): + + doc = Document( + metadata = Metadata( + id = id, + metadata = document.metadata, + user = document.user, + collection = document.collection + ), + data = document.document + ) + + self.document_load.send(None, doc) + + async def load_text(self, id, document): + + doc = TextDocument( + metadata = Metadata( + id = id, + metadata = document.metadata, + user = document.user, + collection = document.collection + ), + text = document.document + ) + + self.text_load.send(None, doc) + + def parse_request(self, v): + + if v.operation is None: + raise RequestError("Null operation") + + if v.operation == "add": + if ( + v.id and v.document and v.document.metadata and + v.document.document and v.document.kind + ): + return partial( + self.librarian.add, + id = v.id, + document = v.document, + ) + else: + raise RequestError("Invalid call") + + raise RequestError("Invalid operation: " + v.operation) + + async def handle(self, msg): + + v = msg.value() + + # Sender-produced ID + + id = msg.properties()["id"] + + print(f"Handling input {id}...", flush=True) + + try: + func = self.parse_request(v) + except RequestError as e: + resp = LibrarianResponse( + error = Error( + type = "request-error", + message = str(e), + ) + ) + await self.send(resp, properties={"id": id}) + return + + try: + resp = await func() + except RequestError as e: + resp = LibrarianResponse( + error = Error( + type = "request-error", + message = str(e), + ) + ) + await self.send(resp, properties={"id": id}) + return + except Exception as e: + print("Exception:", e, flush=True) + resp = LibrarianResponse( + error = Error( + type = "processing-error", + message = "Unhandled error: " + str(e), + ) + ) + await self.send(resp, properties={"id": id}) + return + + print("Send response...", flush=True) + + await self.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, + ) + + parser.add_argument( + '--minio-host', + default=default_minio_host, + help=f'Minio hostname (default: {default_minio_host})', + ) + + parser.add_argument( + '--minio-access-key', + default='minioadmin', + help='Minio access key / username ' + f'(default: {default_minio_access_key})', + ) + + parser.add_argument( + '--minio-secret-key', + default='minioadmin', + help='Minio secret key / password ' + f'(default: {default_minio_access_key})', + ) + + 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' + ) + + 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__) + diff --git a/trustgraph-flow/trustgraph/librarian/table_store.py b/trustgraph-flow/trustgraph/librarian/table_store.py new file mode 100644 index 00000000..99e0d845 --- /dev/null +++ b/trustgraph-flow/trustgraph/librarian/table_store.py @@ -0,0 +1,131 @@ +from .. schema import LibrarianRequest, LibrarianResponse, Error +from .. knowledge import hash +from .. exceptions import RequestError + +from cassandra.cluster import Cluster +from cassandra.auth import PlainTextAuthProvider +from cassandra.query import BatchStatement +import uuid +import time + +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: + auth_provider = PlainTextAuthProvider( + username=cassandra_user, password=cassandra_password + ) + self.cluster = Cluster( + cassandra_host, + auth_provider=auth_provider + ) + else: + self.cluster = Cluster(cassandra_host) + + self.cassandra = self.cluster.connect() + + print("Connected.", flush=True) + + self.ensure_cassandra_schema() + + self.insert_document_stmt = self.cassandra.prepare(""" + insert into document + (id, user, collection, kind, object_id, metadata) + values (?, ?, ?, ?, ?, ?) + """) + + 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("document table...", flush=True) + + self.cassandra.execute(""" + create table if not exists document ( + user text, + collection text, + id uuid, + kind text, + object_id uuid, + metadata list>, + PRIMARY KEY (user, collection, id) + ); + """); + + print("object index...", flush=True) + + self.cassandra.execute(""" + create index if not exists document_object + on document ( object_id) + """); + + print("Cassandra schema OK.", flush=True) + + def add(self, object_id, document): + + if document.kind not in ( + "text/plain", "application/pdf" + ): + raise RequestError("Invalid document kind: " + document.kind) + + # Create random doc ID + doc_id = uuid.uuid4() + + print("Adding", object_id, doc_id) + + 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 document.metadata + ] + + while True: + + try: + + resp = self.cassandra.execute( + self.insert_document_stmt, + ( + doc_id, document.user, document.collection, + document.kind, object_id, metadata + ) + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + time.sleep(1) + + print("Add complete", flush=True) + + + + diff --git a/trustgraph-flow/trustgraph/metering/counter.py b/trustgraph-flow/trustgraph/metering/counter.py index 6e6b829b..68ddf441 100644 --- a/trustgraph-flow/trustgraph/metering/counter.py +++ b/trustgraph-flow/trustgraph/metering/counter.py @@ -57,7 +57,7 @@ class Processor(Consumer): return model["input_price"], model["output_price"] return None, None # Return None if model is not found - def handle(self, msg): + async def handle(self, msg): v = msg.value() modelname = v.model @@ -98,4 +98,4 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/prompt/generic/service.py b/trustgraph-flow/trustgraph/model/prompt/generic/service.py index 89cc850c..8d533fcf 100755 --- a/trustgraph-flow/trustgraph/model/prompt/generic/service.py +++ b/trustgraph-flow/trustgraph/model/prompt/generic/service.py @@ -78,7 +78,7 @@ class Processor(ConsumerProducer): return json.loads(json_str) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -92,32 +92,32 @@ class Processor(ConsumerProducer): if kind == "extract-definitions": - self.handle_extract_definitions(id, v) + await self.handle_extract_definitions(id, v) return elif kind == "extract-topics": - self.handle_extract_topics(id, v) + await self.handle_extract_topics(id, v) return elif kind == "extract-relationships": - self.handle_extract_relationships(id, v) + await self.handle_extract_relationships(id, v) return elif kind == "extract-rows": - self.handle_extract_rows(id, v) + await self.handle_extract_rows(id, v) return elif kind == "kg-prompt": - self.handle_kg_prompt(id, v) + await self.handle_kg_prompt(id, v) return elif kind == "document-prompt": - self.handle_document_prompt(id, v) + await self.handle_document_prompt(id, v) return else: @@ -125,7 +125,7 @@ class Processor(ConsumerProducer): print("Invalid kind.", flush=True) return - def handle_extract_definitions(self, id, v): + async def handle_extract_definitions(self, id, v): try: @@ -164,7 +164,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = PromptResponse(definitions=output, error=None) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -182,9 +182,9 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) - def handle_extract_topics(self, id, v): + async def handle_extract_topics(self, id, v): try: @@ -223,7 +223,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = PromptResponse(topics=output, error=None) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -241,9 +241,9 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) - def handle_extract_relationships(self, id, v): + async def handle_extract_relationships(self, id, v): try: @@ -295,7 +295,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = PromptResponse(relationships=output, error=None) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -313,9 +313,9 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) - def handle_extract_rows(self, id, v): + async def handle_extract_rows(self, id, v): try: @@ -366,7 +366,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = PromptResponse(rows=output, error=None) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -384,9 +384,9 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) - def handle_kg_prompt(self, id, v): + async def handle_kg_prompt(self, id, v): try: @@ -400,7 +400,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = PromptResponse(answer=ans, error=None) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -418,9 +418,9 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) - def handle_document_prompt(self, id, v): + async def handle_document_prompt(self, id, v): try: @@ -437,7 +437,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = PromptResponse(answer=ans, error=None) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) print("Done.", flush=True) @@ -455,7 +455,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) @staticmethod def add_args(parser): @@ -481,5 +481,5 @@ def run(): raise RuntimeError("NOT IMPLEMENTED") - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py index b8aa489c..58657d7d 100755 --- a/trustgraph-flow/trustgraph/model/prompt/template/service.py +++ b/trustgraph-flow/trustgraph/model/prompt/template/service.py @@ -156,7 +156,7 @@ class Processor(ConsumerProducer): config = prompt_configuration, ) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -191,7 +191,7 @@ class Processor(ConsumerProducer): error=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) return @@ -206,7 +206,7 @@ class Processor(ConsumerProducer): error=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) return @@ -224,7 +224,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) except Exception as e: @@ -240,7 +240,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) @staticmethod def add_args(parser): @@ -294,5 +294,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py index 90be6962..8951b71e 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py @@ -123,7 +123,7 @@ class Processor(ConsumerProducer): return result - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -154,7 +154,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) except TooManyRequests: @@ -182,7 +182,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -224,4 +224,4 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py index f5ecb8d6..53833898 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -84,7 +84,7 @@ class Processor(ConsumerProducer): azure_endpoint = endpoint, ) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -133,7 +133,7 @@ class Processor(ConsumerProducer): model=self.model ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) except RateLimitError: @@ -161,7 +161,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -212,4 +212,4 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py index 5cfd8907..195a39e4 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -73,7 +73,7 @@ class Processor(ConsumerProducer): print("Initialised", flush=True) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -150,7 +150,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -190,6 +190,6 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index 5b8e3ba9..4b18f7f1 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -69,7 +69,7 @@ class Processor(ConsumerProducer): print("Initialised", flush=True) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -106,7 +106,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) - self.send(r, properties={"id": id}) + self.await send(r, properties={"id": id}) print("Done.", flush=True) @@ -136,7 +136,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.producer.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -169,6 +169,6 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index 5d5b23a0..98ecaf0e 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -102,7 +102,7 @@ class Processor(ConsumerProducer): print("Initialised", flush=True) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -146,7 +146,7 @@ class Processor(ConsumerProducer): out_token=outputtokens, model=self.model ) - self.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -177,7 +177,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -217,6 +217,6 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py index 65a2b171..483412a2 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py @@ -74,7 +74,7 @@ class Processor(ConsumerProducer): print("Initialised", flush=True) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -122,7 +122,7 @@ class Processor(ConsumerProducer): out_token=outputtokens, model="llama.cpp" ) - self.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -145,7 +145,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -185,6 +185,6 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py index 8c5bd3dc..6d825bac 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py @@ -71,7 +71,7 @@ class Processor(ConsumerProducer): self.model = model self.llm = Client(host=ollama) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -96,7 +96,7 @@ class Processor(ConsumerProducer): r = TextCompletionResponse(response=response_text, error=None, in_token=inputtokens, out_token=outputtokens, model="ollama") - self.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -119,7 +119,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -145,6 +145,6 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index c2b948d5..ebfae9ed 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py @@ -73,7 +73,7 @@ class Processor(ConsumerProducer): print("Initialised", flush=True) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -126,7 +126,7 @@ class Processor(ConsumerProducer): out_token=outputtokens, model=self.model ) - self.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -156,7 +156,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -196,6 +196,6 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/milvus/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/milvus/service.py index 8e106e6f..b16399e9 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/milvus/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/milvus/service.py @@ -40,7 +40,7 @@ class Processor(ConsumerProducer): self.vecstore = DocVectors(store_uri) - def handle(self, msg): + async def handle(self, msg): try: @@ -64,7 +64,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = DocumentEmbeddingsResponse(documents=chunks, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -82,7 +82,7 @@ class Processor(ConsumerProducer): documents=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -102,5 +102,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py index b8502143..6a88671c 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py @@ -56,7 +56,7 @@ class Processor(ConsumerProducer): } ) - def handle(self, msg): + async def handle(self, msg): try: @@ -100,7 +100,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = DocumentEmbeddingsResponse(documents=chunks, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -118,7 +118,7 @@ class Processor(ConsumerProducer): documents=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -143,5 +143,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py index dd53862f..128203ad 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py @@ -47,7 +47,7 @@ class Processor(ConsumerProducer): self.client = QdrantClient(url=store_uri, api_key=api_key) - def handle(self, msg): + async def handle(self, msg): try: @@ -81,7 +81,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = DocumentEmbeddingsResponse(documents=chunks, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -99,7 +99,7 @@ class Processor(ConsumerProducer): documents=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -125,5 +125,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py index b5f9ae5b..8dd8d04d 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py @@ -46,7 +46,7 @@ class Processor(ConsumerProducer): else: return Value(value=ent, is_uri=False) - def handle(self, msg): + async def handle(self, msg): try: @@ -79,7 +79,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = GraphEmbeddingsResponse(entities=entities, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -97,7 +97,7 @@ class Processor(ConsumerProducer): entities=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -117,5 +117,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py index 2534d278..90cfc6de 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py @@ -62,7 +62,7 @@ class Processor(ConsumerProducer): else: return Value(value=ent, is_uri=False) - def handle(self, msg): + async def handle(self, msg): try: @@ -120,7 +120,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = GraphEmbeddingsResponse(entities=entities, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -138,7 +138,7 @@ class Processor(ConsumerProducer): entities=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -163,5 +163,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index 4e9492c6..dc3e28f3 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py @@ -52,7 +52,7 @@ class Processor(ConsumerProducer): else: return Value(value=ent, is_uri=False) - def handle(self, msg): + async def handle(self, msg): try: @@ -106,7 +106,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = GraphEmbeddingsResponse(entities=entities, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -124,7 +124,7 @@ class Processor(ConsumerProducer): entities=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -150,5 +150,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py index 22fbf84d..e3687756 100755 --- a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py +++ b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py @@ -53,7 +53,7 @@ class Processor(ConsumerProducer): else: return Value(value=ent, is_uri=False) - def handle(self, msg): + async def handle(self, msg): try: @@ -154,7 +154,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = TriplesQueryResponse(triples=triples, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -172,7 +172,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -205,5 +205,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/triples/falkordb/service.py b/trustgraph-flow/trustgraph/query/triples/falkordb/service.py index 1d77bb15..56fed6d3 100755 --- a/trustgraph-flow/trustgraph/query/triples/falkordb/service.py +++ b/trustgraph-flow/trustgraph/query/triples/falkordb/service.py @@ -54,7 +54,7 @@ class Processor(ConsumerProducer): else: return Value(value=ent, is_uri=False) - def handle(self, msg): + async def handle(self, msg): try: @@ -301,7 +301,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = TriplesQueryResponse(triples=triples, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -319,7 +319,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -345,5 +345,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/triples/memgraph/service.py b/trustgraph-flow/trustgraph/query/triples/memgraph/service.py index 46dd19e3..f442c4ef 100755 --- a/trustgraph-flow/trustgraph/query/triples/memgraph/service.py +++ b/trustgraph-flow/trustgraph/query/triples/memgraph/service.py @@ -58,7 +58,7 @@ class Processor(ConsumerProducer): else: return Value(value=ent, is_uri=False) - def handle(self, msg): + async def handle(self, msg): try: @@ -313,7 +313,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = TriplesQueryResponse(triples=triples, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -331,7 +331,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -369,5 +369,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/query/triples/neo4j/service.py b/trustgraph-flow/trustgraph/query/triples/neo4j/service.py index d60bc4f4..49ba0345 100755 --- a/trustgraph-flow/trustgraph/query/triples/neo4j/service.py +++ b/trustgraph-flow/trustgraph/query/triples/neo4j/service.py @@ -58,7 +58,7 @@ class Processor(ConsumerProducer): else: return Value(value=ent, is_uri=False) - def handle(self, msg): + async def handle(self, msg): try: @@ -297,7 +297,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = TriplesQueryResponse(triples=triples, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -315,7 +315,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -353,5 +353,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py index 0d2a3b86..23b46129 100755 --- a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py @@ -79,7 +79,7 @@ class Processor(ConsumerProducer): module=module, ) - def handle(self, msg): + async def handle(self, msg): try: @@ -94,7 +94,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = DocumentRagResponse(response = response, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -112,7 +112,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -162,5 +162,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index 5d5793d6..4095d0c3 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py @@ -98,7 +98,7 @@ class Processor(ConsumerProducer): module=module, ) - def handle(self, msg): + async def handle(self, msg): try: @@ -115,7 +115,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = GraphRagResponse(response = response, error=None) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -133,7 +133,7 @@ class Processor(ConsumerProducer): response=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) @@ -216,5 +216,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py index bfa6c123..b4dbc486 100755 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py @@ -35,7 +35,7 @@ class Processor(Consumer): self.vecstore = DocVectors(store_uri) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -65,5 +65,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py index c59ecd7b..9e91db9a 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py @@ -61,7 +61,7 @@ class Processor(Consumer): self.last_index_name = None - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -166,5 +166,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index a01fd9e4..810c1931 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -42,7 +42,7 @@ class Processor(Consumer): self.client = QdrantClient(url=store_uri) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -110,5 +110,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py index e1379577..b2d40306 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py @@ -34,7 +34,7 @@ class Processor(Consumer): self.vecstore = EntityVectors(store_uri) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -59,5 +59,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py index a32ff627..83861b54 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py @@ -88,7 +88,7 @@ class Processor(Consumer): "Gave up waiting for index creation" ) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -170,5 +170,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index ffe33565..6b0d7371 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -67,7 +67,7 @@ class Processor(Consumer): return cname - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -117,5 +117,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/object_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/object_embeddings/milvus/write.py index 468b357a..5490af97 100755 --- a/trustgraph-flow/trustgraph/storage/object_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/object_embeddings/milvus/write.py @@ -34,7 +34,7 @@ class Processor(Consumer): self.vecstore = ObjectVectors(store_uri) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -57,5 +57,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py b/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py index fc8f6686..c8f3b9e1 100755 --- a/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py @@ -62,7 +62,7 @@ class Processor(Consumer): self.session.execute("use trustgraph"); - def handle(self, msg): + async def handle(self, msg): try: @@ -143,5 +143,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py index d940d0ec..17b5ae9a 100755 --- a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py @@ -47,7 +47,7 @@ class Processor(Consumer): self.password = graph_password self.table = None - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -110,5 +110,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/triples/falkordb/write.py b/trustgraph-flow/trustgraph/storage/triples/falkordb/write.py index 3c7d1660..2d0ae38a 100755 --- a/trustgraph-flow/trustgraph/storage/triples/falkordb/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/falkordb/write.py @@ -118,7 +118,7 @@ class Processor(Consumer): time=res.run_time_ms )) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -154,5 +154,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py b/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py index 7295e691..620e669e 100755 --- a/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py @@ -205,7 +205,7 @@ class Processor(Consumer): src=t.s.value, dest=t.o.value, uri=t.p.value, ) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -256,5 +256,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py index 18b40129..3323f912 100755 --- a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py @@ -158,7 +158,7 @@ class Processor(Consumer): time=summary.result_available_after )) - def handle(self, msg): + async def handle(self, msg): v = msg.value() @@ -206,5 +206,5 @@ class Processor(Consumer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__) diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index d6a2efec..efe63ea6 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -131,7 +131,7 @@ class Processor(ConsumerProducer): print("Initialisation complete", flush=True) - def handle(self, msg): + async def handle(self, msg): try: @@ -248,5 +248,5 @@ class Processor(ConsumerProducer): def run(): - Processor.start(module, __doc__) + Processor.launch(module, __doc__)