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,