From e2896f64337ab657d43faf9ecb95583da20a0655 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 12 Aug 2024 17:14:02 +0100 Subject: [PATCH] Add graph embeddings query --- scripts/ge-query-milvus | 6 ++ setup.py | 1 + trustgraph/query/__init__.py | 0 trustgraph/query/graph_embeddings/__init__.py | 0 .../query/graph_embeddings/milvus/__init__.py | 3 + .../query/graph_embeddings/milvus/__main__.py | 7 ++ .../query/graph_embeddings/milvus/service.py | 83 +++++++++++++++++++ trustgraph/schema.py | 3 +- .../storage/graph_embeddings/milvus/write.py | 4 +- 9 files changed, 104 insertions(+), 3 deletions(-) create mode 100755 scripts/ge-query-milvus create mode 100644 trustgraph/query/__init__.py create mode 100644 trustgraph/query/graph_embeddings/__init__.py create mode 100644 trustgraph/query/graph_embeddings/milvus/__init__.py create mode 100755 trustgraph/query/graph_embeddings/milvus/__main__.py create mode 100755 trustgraph/query/graph_embeddings/milvus/service.py diff --git a/scripts/ge-query-milvus b/scripts/ge-query-milvus new file mode 100755 index 00000000..179750cb --- /dev/null +++ b/scripts/ge-query-milvus @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.query.graph_embeddings.milvus import run + +run() + diff --git a/setup.py b/setup.py index 1230d377..4649f9bb 100644 --- a/setup.py +++ b/setup.py @@ -57,6 +57,7 @@ setuptools.setup( "scripts/embeddings-ollama", "scripts/embeddings-vectorize", "scripts/ge-dump-parquet", + "scripts/ge-query-milvus", "scripts/ge-write-milvus", "scripts/graph-rag", "scripts/graph-show", diff --git a/trustgraph/query/__init__.py b/trustgraph/query/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/trustgraph/query/graph_embeddings/__init__.py b/trustgraph/query/graph_embeddings/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/trustgraph/query/graph_embeddings/milvus/__init__.py b/trustgraph/query/graph_embeddings/milvus/__init__.py new file mode 100644 index 00000000..ba844705 --- /dev/null +++ b/trustgraph/query/graph_embeddings/milvus/__init__.py @@ -0,0 +1,3 @@ + +from . service import * + diff --git a/trustgraph/query/graph_embeddings/milvus/__main__.py b/trustgraph/query/graph_embeddings/milvus/__main__.py new file mode 100755 index 00000000..89684e3e --- /dev/null +++ b/trustgraph/query/graph_embeddings/milvus/__main__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from . hf import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph/query/graph_embeddings/milvus/service.py b/trustgraph/query/graph_embeddings/milvus/service.py new file mode 100755 index 00000000..456ae587 --- /dev/null +++ b/trustgraph/query/graph_embeddings/milvus/service.py @@ -0,0 +1,83 @@ + +""" +Graph embeddings query service. Input is vector, output is list of +e +""" + +from ... direct.milvus import TripleVectors +from ... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse +from ... schema import graph_embeddings_request_queue +from ... schema import graph_embeddings_response_queue +from ... base import ConsumerProducer + +module = ".".join(__name__.split(".")[1:-1]) + +default_input_queue = graph_embeddings_request_queue +default_output_queue = graph_embeddings_response_queue +default_subscriber = module +default_store_uri = 'http://localhost:19530' + +class Processor(ConsumerProducer): + + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + store_uri = params.get("store_uri", default_store_uri) + + super(Processor, self).__init__( + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": GraphEmbeddingsRequest, + "output_schema": GraphEmbeddingsResponse, + "store_uri": store_uri, + } + ) + + self.vecstore = TripleVectors(store_uri) + + def handle(self, msg): + + v = msg.value() + + # Sender-produced ID + id = msg.properties()["id"] + + print(f"Handling input {id}...", flush=True) + + entities = set() + + for vec in v.vectors: + + resp = self.vecstore.search(vec, limit=v.limit) + + for r in resp: + entities.add(r) + + print("Send response...", flush=True) + r = GraphEmbeddingsResponse(entities=entities) + self.producer.send(r, 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( + '-t', '--store-uri', + default=default_store_uri, + help=f'Milvus store URI (default: {default_store_uri})' + ) + +def run(): + + Processor.start(module, __doc__) + diff --git a/trustgraph/schema.py b/trustgraph/schema.py index d82a7e20..f0dbb8ff 100644 --- a/trustgraph/schema.py +++ b/trustgraph/schema.py @@ -75,9 +75,10 @@ graph_embeddings_store_queue = topic('graph-embeddings-store') class GraphEmbeddingsQueryRequest(Record): vectors = Array(Array(Double())) + limit = Integer() class GraphEmbeddingsQueryResponse(Record): - vectors = Array(Value()) + entities = Array(Value()) graph_embeddings_request_queue = topic( 'graph-embeddings', kind='non-persistent', namespace='request' diff --git a/trustgraph/storage/graph_embeddings/milvus/write.py b/trustgraph/storage/graph_embeddings/milvus/write.py index e5141ca8..9256f011 100755 --- a/trustgraph/storage/graph_embeddings/milvus/write.py +++ b/trustgraph/storage/graph_embeddings/milvus/write.py @@ -51,8 +51,8 @@ class Processor(Consumer): parser.add_argument( '-t', '--store-uri', - default="http://milvus:19530", - help=f'Milvus store URI (default: http://milvus:19530)' + default=default_store_uri, + help=f'Milvus store URI (default: {default_store_uri})' ) def run():