Query side starting to come together

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 \

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__)
+

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__)
+

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__)
 

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(