diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 61e75ace..bc40d156 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/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 diff --git a/trustgraph-base/trustgraph/base/embeddings_service.py b/trustgraph-base/trustgraph/base/embeddings_service.py new file mode 100755 index 00000000..c6befdb7 --- /dev/null +++ b/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) + + + diff --git a/trustgraph-base/trustgraph/base/llm_service.py b/trustgraph-base/trustgraph/base/llm_service.py index 096ba224..d22fe564 100755 --- a/trustgraph-base/trustgraph/base/llm_service.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py index 46716364..10e3e62a 100755 --- a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py +++ b/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__) + diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py index 61606405..cf8e2d8a 100755 --- a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py +++ b/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__)