Graph embeddings works

trustgraph-flow/trustgraph/embeddings/fastembed/processor.py

@@ -1,55 +1,42 @@
 
 """
-Embeddings service, applies an embeddings model selected from HuggingFace.
+Embeddings service, applies an embeddings model using fastembed
 Input is text, output is embeddings vector.
 """
 
-from ... schema import EmbeddingsRequest, EmbeddingsResponse
-from ... base import RequestResponseService
+from ... base import EmbeddingsService
 
 from fastembed import TextEmbedding
-import os
 
-module = "embeddings"
-
-default_subscriber = module
+default_ident = "embeddings"
 default_model="sentence-transformers/all-MiniLM-L6-v2"
 
-class Processor(RequestResponseService):
+class Processor(EmbeddingsService):
 
     def __init__(self, **params):
 
         model = params.get("model", default_model)
 
         super(Processor, self).__init__(
-            **params | {
-                "model": model,
-                "request_schema": EmbeddingsRequest,
-                "response_schema": EmbeddingsResponse,
-            }
+            **params | { "model": model }
         )
 
+        print("Get model...", flush=True)
         self.embeddings = TextEmbedding(model_name = model)
 
-    async def on_request(self, request, consumer, flow):
+    async def on_embeddings(self, text):
 
-        text = request.text
         vecs = self.embeddings.embed([text])
 
-        vecs = [
+        return [
             v.tolist()
             for v in vecs
         ]
 
-        return EmbeddingsResponse(
-            vectors=list(vecs),
-            error=None,
-        )
-
     @staticmethod
     def add_args(parser):
 
-        RequestResponseService.add_args(parser, default_subscriber)
+        EmbeddingsService.add_args(parser)
 
         parser.add_argument(
             '-m', '--model',
@@ -59,5 +46,6 @@ class Processor(RequestResponseService):
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)
+
 

trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py

@@ -6,58 +6,50 @@ Output is entity plus embedding.
 """
 
 from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings
-from ... schema import entity_contexts_ingest_queue
-from ... schema import graph_embeddings_store_queue
-from ... schema import embeddings_request_queue, embeddings_response_queue
-from ... clients.embeddings_client import EmbeddingsClient
-from ... log_level import LogLevel
-from ... base import FlowProcessor
+from ... schema import EmbeddingsRequest, EmbeddingsResponse
 
-module = "graph-embeddings"
+from ... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
+from ... base import ProducerSpec
 
-default_subscriber = module
+default_ident = "graph-embeddings"
 
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
-
-
-                "input_schema": EntityContexts,
-                "output_schema": GraphEmbeddings,
-        
-
-
-
         id = params.get("id")
-        subscriber = params.get("subscriber", default_subscriber)
 
         super(Processor, self).__init__(
             **params | {
                 "id": id,
-                "subscriber": subscriber,
             }
         )
 
-        self.register_consumer(
-            name = "input",
-            schema = EntityContexts,
-            handler = self.on_message,
+        self.register_specification(
+            ConsumerSpec(
+                name = "input",
+                schema = EntityContexts,
+                handler = self.on_message,
+            )
         )
 
-        self.register_producer(
-            name = "output",
-            schema = GraphEmbeddings,
+        self.register_specification(
+            RequestResponseSpec(
+                request_name = "embeddings-request",
+                request_schema = EmbeddingsRequest,
+                response_name = "embeddings-response",
+                response_schema = EmbeddingsResponse,
+            )
         )
 
-        # self.embeddings = EmbeddingsClient(
-        #     pulsar_host=self.pulsar_host,
-        #     input_queue=emb_request_queue,
-        #     output_queue=emb_response_queue,
-        #     subscriber=module + "-emb",
-        # )
+        self.register_specification(
+            ProducerSpec(
+                name = "output",
+                schema = GraphEmbeddings
+            )
+        )
 
-    async def handle(self, msg):
+    async def on_message(self, msg, consumer, flow):
 
         v = msg.value()
         print(f"Indexing {v.metadata.id}...", flush=True)
@@ -68,7 +60,13 @@ class Processor(FlowProcessor):
 
             for entity in v.entities:
 
-                vectors = self.embeddings.request(entity.context)
+                resp = await flow("embeddings-request").request(
+                    EmbeddingsRequest(
+                        text = entity.context
+                    )
+                )
+
+                vectors = resp.vectors
 
                 entities.append(
                     EntityEmbeddings(
@@ -82,7 +80,7 @@ class Processor(FlowProcessor):
                 entities=entities,
             )
 
-            await self.send(r)
+            await flow("output").send(r)
 
         except Exception as e:
             print("Exception:", e, flush=True)
@@ -95,24 +93,9 @@ class Processor(FlowProcessor):
     @staticmethod
     def add_args(parser):
 
-        ConsumerProducer.add_args(
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
-
-        parser.add_argument(
-            '--embeddings-request-queue',
-            default=embeddings_request_queue,
-            help=f'Embeddings request queue (default: {embeddings_request_queue})',
-        )
-
-        parser.add_argument(
-            '--embeddings-response-queue',
-            default=embeddings_response_queue,
-            help=f'Embeddings request queue (default: {embeddings_response_queue})',
-        )
+        FlowProcessor.add_args(parser)
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)