diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 9086fd8c..ddf47f97 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -18,4 +18,6 @@ from . embeddings_client import EmbeddingsClientSpec from . text_completion_client import TextCompletionClientSpec from . prompt_client import PromptClientSpec from . triples_store_service import TriplesStoreService +from . graph_embeddings_store_service import GraphEmbeddingsStoreService +from . document_embeddings_store_service import DocumentEmbeddingsStoreService diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 4427a09d..9bf5f33c 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -8,31 +8,21 @@ from qdrant_client.models import PointStruct from qdrant_client.models import Distance, VectorParams import uuid -from .... schema import GraphEmbeddings -from .... schema import graph_embeddings_store_queue -from .... log_level import LogLevel -from .... base import Consumer +from .... base import GraphEmbeddingsStoreService -module = "ge-write" +default_ident = "ge-write" -default_input_queue = graph_embeddings_store_queue -default_subscriber = module default_store_uri = 'http://localhost:6333' -class Processor(Consumer): +class Processor(GraphEmbeddingsStoreService): def __init__(self, **params): - input_queue = params.get("input_queue", default_input_queue) - subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { - "input_queue": input_queue, - "subscriber": subscriber, - "input_schema": GraphEmbeddings, "store_uri": store_uri, "api_key": api_key, } @@ -67,11 +57,9 @@ class Processor(Consumer): return cname - async def handle(self, msg): + async def store_graph_embeddings(self, message): - v = msg.value() - - for entity in v.entities: + for entity in message.entities: if entity.entity.value == "" or entity.entity.value is None: return @@ -99,9 +87,7 @@ class Processor(Consumer): @staticmethod def add_args(parser): - Consumer.add_args( - parser, default_input_queue, default_subscriber, - ) + GraphEmbeddingsStoreService.add_args(parser) parser.add_argument( '-t', '--store-uri', @@ -117,5 +103,5 @@ class Processor(Consumer): def run(): - Processor.launch(module, __doc__) + Processor.launch(default_ident, __doc__)