Graph embeddings works

trustgraph-base/trustgraph/base/__init__.py

@@ -13,6 +13,6 @@ from . producer_spec import ProducerSpec
 from . subscriber_spec import SubscriberSpec
 from . request_response_spec import RequestResponseSpec
 from . llm_service import LlmService, LlmResult
-
+from . embeddings_service import EmbeddingsService
 
 

trustgraph-base/trustgraph/base/embeddings_service.py

@@ -0,0 +1,90 @@
+
+"""
+Embeddings resolution base class
+"""
+
+import time
+from prometheus_client import Histogram
+
+from .. schema import EmbeddingsRequest, EmbeddingsResponse, Error
+from .. exceptions import TooManyRequests
+from .. base import FlowProcessor, ConsumerSpec, ProducerSpec
+
+default_ident = "embeddings"
+
+class EmbeddingsService(FlowProcessor):
+
+    def __init__(self, **params):
+
+        id = params.get("id")
+
+        super(EmbeddingsService, self).__init__(**params | { "id": id })
+
+        self.register_specification(
+            ConsumerSpec(
+                name = "request",
+                schema = EmbeddingsRequest,
+                handler = self.on_request
+            )
+        )
+
+        self.register_specification(
+            ProducerSpec(
+                name = "response",
+                schema = EmbeddingsResponse
+            )
+        )
+
+    async def on_request(self, msg, consumer, flow):
+
+        try:
+
+            request = msg.value()
+
+            # Sender-produced ID
+
+            id = msg.properties()["id"]
+
+            print("Handling request", id, "...", flush=True)
+
+            vectors = await self.on_embeddings(request.text)
+
+            await flow("response").send(
+                EmbeddingsResponse(
+                    error = None,
+                    vectors = vectors,
+                ),
+                properties={"id": id}
+            )
+
+            print("Handled.", flush=True)
+
+        except TooManyRequests as e:
+            raise e
+
+        except Exception as e:
+
+            # Apart from rate limits, treat all exceptions as unrecoverable
+
+            print(f"Exception: {e}", flush=True)
+
+            print("Send error response...", flush=True)
+
+            await flow.producer["response"].send(
+                EmbeddingsResponse(
+                    error=Error(
+                        type = "embeddings-error",
+                        message = str(e),
+                    ),
+                    vectors=None,
+                ),
+                properties={"id": id}
+            )
+
+    @staticmethod
+    def add_args(parser):
+
+        FlowProcessor.add_args(parser)
+
+
+

trustgraph-base/trustgraph/base/llm_service.py

@@ -73,7 +73,7 @@ class LlmService(FlowProcessor):
                     request.system, request.prompt
                 )
 
-            await flow.producer["response"].send(
+            await flow("response").send(
                 TextCompletionResponse(
                     error=None,
                     response=response.text,

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