diff --git a/Makefile b/Makefile index 1899e602..73edcf2f 100644 --- a/Makefile +++ b/Makefile @@ -65,8 +65,8 @@ some-containers: -t ${CONTAINER_BASE}/trustgraph-base:${VERSION} . ${DOCKER} build -f containers/Containerfile.flow \ -t ${CONTAINER_BASE}/trustgraph-flow:${VERSION} . - ${DOCKER} build -f containers/Containerfile.vertexai \ - -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} . +# ${DOCKER} build -f containers/Containerfile.vertexai \ +# -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} . basic-containers: update-package-versions ${DOCKER} build -f containers/Containerfile.base \ diff --git a/trustgraph-base/trustgraph/base/graph_embeddings_query_service.py b/trustgraph-base/trustgraph/base/graph_embeddings_query_service.py new file mode 100755 index 00000000..04efed0c --- /dev/null +++ b/trustgraph-base/trustgraph/base/graph_embeddings_query_service.py @@ -0,0 +1,83 @@ + +""" +Graph embeddings query service. Input is vectors. Output is list of +embeddings. +""" + +from .... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse +from .... schema import Error, Value + +from .... base import FlowProcessor, GraphEmbeddingsQueryService, ConsumerSpec +from .... base import ProducerSpec + +default_ident = "ge-query" + +class GraphEmbeddingsQueryService(FlowProcessor): + + def __init__(self, **params): + + id = params.get("id") + + super(GraphEmbeddingsQueryService, self).__init__( + **params | { "id": id } + ) + + self.register_specification( + ConsumerSpec( + name = "request", + schema = GraphEmbeddingsRequest, + handler = self.on_message + ) + ) + + self.register_specification( + ProducerSpec( + name = "response", + schema = GraphEmbeddingsResponse, + ) + ) + + async def on_message(self, msg, consumer, flow): + + try: + + request = msg.value() + + # Sender-produced ID + id = msg.properties()["id"] + + print(f"Handling input {id}...", flush=True) + + entities = self.query_graph_embeddings(request) + + print("Send response...", flush=True) + r = GraphEmbeddingsResponse(entities=entities, error=None) + await flow("response").send(r, properties={"id": id}) + + print("Done.", flush=True) + + except Exception as e: + + print(f"Exception: {e}") + + print("Send error response...", flush=True) + + r = GraphEmbeddingsResponse( + error=Error( + type = "graph-embeddings-query-error", + message = str(e), + ), + response=None, + ) + + await flow("response").send(r, properties={"id": id}) + + @staticmethod + def add_args(parser): + + GraphEmbeddingsQueryService.add_args(parser) + +def run(): + + Processor.launch(default_ident, __doc__) + diff --git a/trustgraph-base/trustgraph/base/triples_query_service.py b/trustgraph-base/trustgraph/base/triples_query_service.py new file mode 100755 index 00000000..0538a7f7 --- /dev/null +++ b/trustgraph-base/trustgraph/base/triples_query_service.py @@ -0,0 +1,81 @@ + +""" +Triples query service. Input is a (s, p, o) triple, some values may be +null. Output is a list of triples. +""" + +from .... schema import TriplesQueryRequest, TriplesQueryResponse, Error +from .... schema import Value, Triple + +from .... base import FlowProcessor, TriplesQueryService, ConsumerSpec +from .... base import ProducerSpec + +default_ident = "triples-query" + +class TriplesQueryService(FlowProcessor): + + def __init__(self, **params): + + id = params.get("id") + + super(TriplesStoreService, self).__init__(**params | { "id": id }) + + self.register_specification( + ConsumerSpec( + name = "request", + schema = TriplesQueryRequest, + handler = self.on_message + ) + ) + + self.register_specification( + ProducerSpec( + name = "response", + schema = TriplesQueryResponse, + ) + ) + + async def on_message(self, msg, consumer, flow): + + try: + + request = msg.value() + + # Sender-produced ID + id = msg.properties()["id"] + + print(f"Handling input {id}...", flush=True) + + triples = self.query_triples(request) + + print("Send response...", flush=True) + r = TriplesQueryResponse(triples=triples, error=None) + await flow("response").send(r, properties={"id": id}) + + print("Done.", flush=True) + + except Exception as e: + + print(f"Exception: {e}") + + print("Send error response...", flush=True) + + r = TriplesQueryResponse( + error = Error( + type = "triples-query-error", + message = str(e), + ), + triples = None, + ) + + await flow("response").send(r, properties={"id": id}) + + @staticmethod + def add_args(parser): + + FlowProcessor.add_args(parser) + +def run(): + + Processor.launch(default_ident, __doc__) + diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index 28a1637c..74a5d182 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py @@ -7,44 +7,28 @@ entities from qdrant_client import QdrantClient from qdrant_client.models import PointStruct from qdrant_client.models import Distance, VectorParams -import uuid -from .... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse +from .... schema import GraphEmbeddingsResponse from .... schema import Error, Value -from .... schema import graph_embeddings_request_queue -from .... schema import graph_embeddings_response_queue -from .... base import ConsumerProducer +from .... base import GraphEmbeddingsQueryService -module = "ge-query" - -default_input_queue = graph_embeddings_request_queue -default_output_queue = graph_embeddings_response_queue -default_subscriber = module default_store_uri = 'http://localhost:6333' -class Processor(ConsumerProducer): +class Processor(GraphEmbeddingsQueryService): 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) api_key = params.get("api_key", None) 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, "api_key": api_key, } ) - self.client = QdrantClient(url=store_uri, api_key=api_key) + self.qdrant = QdrantClient(url=store_uri, api_key=api_key) def create_value(self, ent): if ent.startswith("http://") or ent.startswith("https://"): @@ -76,7 +60,7 @@ class Processor(ConsumerProducer): # Heuristic hack, get (2*limit), so that we have more chance # of getting (limit) entities - search_result = self.client.query_points( + search_result = self.qdrant.query_points( collection_name=collection, query=vec, limit=v.limit * 2, @@ -105,36 +89,19 @@ class Processor(ConsumerProducer): entities = ents2 print("Send response...", flush=True) - r = GraphEmbeddingsResponse(entities=entities, error=None) - await self.send(r, properties={"id": id}) + return entities print("Done.", flush=True) except Exception as e: print(f"Exception: {e}") - - print("Send error response...", flush=True) - - r = GraphEmbeddingsResponse( - error=Error( - type = "llm-error", - message = str(e), - ), - entities=None, - ) - - await self.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + raise e @staticmethod def add_args(parser): - ConsumerProducer.add_args( - parser, default_input_queue, default_subscriber, - default_output_queue, - ) + GraphEmbeddingsQueryService.add_args(parser) parser.add_argument( '-t', '--store-uri', @@ -150,5 +117,5 @@ class Processor(ConsumerProducer): def run(): - Processor.launch(module, __doc__) + Processor.launch(default_ident, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 9bf5f33c..d8645f0a 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -30,7 +30,7 @@ class Processor(GraphEmbeddingsStoreService): self.last_collection = None - self.client = QdrantClient(url=store_uri, api_key=api_key) + self.qdrant = QdrantClient(url=store_uri, api_key=api_key) def get_collection(self, dim, user, collection): @@ -40,10 +40,10 @@ class Processor(GraphEmbeddingsStoreService): if cname != self.last_collection: - if not self.client.collection_exists(cname): + if not self.qdrant.collection_exists(cname): try: - self.client.create_collection( + self.qdrant.create_collection( collection_name=cname, vectors_config=VectorParams( size=dim, distance=Distance.COSINE @@ -71,7 +71,7 @@ class Processor(GraphEmbeddingsStoreService): dim, v.metadata.user, v.metadata.collection ) - self.client.upsert( + self.qdrant.upsert( collection_name=collection, points=[ PointStruct(