- Change flawed _client timeout logic which was causing major lags

- Moved clients to trustgraph.clients to tidy the parent directory
- Version bump

trustgraph/base/consumer_producer.py

@@ -126,52 +126,3 @@ class ConsumerProducer(BaseProcessor):
             help=f'Output queue (default: {default_output_queue})'
         )
 
-class Producer(BaseProcessor):
-
-    def __init__(self, **params):
-
-        output_queue = params.get("output_queue")
-        output_schema = params.get("output_schema")
-
-        if not hasattr(__class__, "output_metric"):
-            __class__.output_metric = Counter(
-                'output_count', 'Output items created'
-            )
-
-        if not hasattr(__class__, "pubsub_metric"):
-            __class__.pubsub_metric = Info(
-                'pubsub', 'Pub/sub configuration'
-            )
-
-        __class__.pubsub_metric.info({
-            "output_queue": output_queue,
-            "output_schema": output_schema.__name__,
-        })
-
-        super(Producer, self).__init__(**params)
-
-        if output_schema == None:
-            raise RuntimeError("output_schema must be specified")
-
-        self.producer = self.client.create_producer(
-            topic=output_queue,
-            schema=JsonSchema(output_schema),
-        )
-
-    def send(self, msg, properties={}):
-        self.producer.send(msg, properties)
-        __class__.output_metric.inc()
-
-    @staticmethod
-    def add_args(
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-    ):
-
-        BaseProcessor.add_args(parser)
-
-        parser.add_argument(
-            '-o', '--output-queue',
-            default=default_output_queue,
-            help=f'Output queue (default: {default_output_queue})'
-        )

trustgraph/clients/embeddings_client.py

@@ -1,12 +1,12 @@
-#!/usr/bin/env python3
 
 import pulsar
 import _pulsar
 from pulsar.schema import JsonSchema
-from . schema import EmbeddingsRequest, EmbeddingsResponse
-from . schema import embeddings_request_queue, embeddings_response_queue
+from .. schema import EmbeddingsRequest, EmbeddingsResponse
+from .. schema import embeddings_request_queue, embeddings_response_queue
 import hashlib
 import uuid
+import time
 
 # Ugly
 ERROR=_pulsar.LoggerLevel.Error
@@ -60,9 +60,14 @@ class EmbeddingsClient:
         )
         self.producer.send(r, properties={ "id": id })
 
-        while True:
+        end_time = time.time() + timeout
 
-            msg = self.consumer.receive(timeout_millis=timeout * 1000)
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
 
             mid = msg.properties()["id"]
 
@@ -74,6 +79,8 @@ class EmbeddingsClient:
             # Ignore messages with wrong ID
             self.consumer.acknowledge(msg)
 
+        raise TimeoutError("Timed out waiting for response")
+
     def __del__(self):
 
         if hasattr(self, "consumer"):

trustgraph/clients/graph_embeddings_client.py

@@ -1,14 +1,14 @@
-#!/usr/bin/env python3
 
 import pulsar
 import _pulsar
 from pulsar.schema import JsonSchema
 import hashlib
 import uuid
+import time
 
-from . schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
-from . schema import graph_embeddings_request_queue
-from . schema import graph_embeddings_response_queue
+from .. schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
+from .. schema import graph_embeddings_request_queue
+from .. schema import graph_embeddings_response_queue
 
 # Ugly
 ERROR=_pulsar.LoggerLevel.Error
@@ -62,9 +62,14 @@ class GraphEmbeddingsClient:
 
         self.producer.send(r, properties={ "id": id })
 
-        while True:
+        end_time = time.time() + timeout
 
-            msg = self.consumer.receive(timeout_millis=timeout * 1000)
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
 
             mid = msg.properties()["id"]
 
@@ -76,6 +81,8 @@ class GraphEmbeddingsClient:
             # Ignore messages with wrong ID
             self.consumer.acknowledge(msg)
 
+        raise TimeoutError("Timed out waiting for response")
+
     def __del__(self):
 
         if hasattr(self, "consumer"):

trustgraph/clients/graph_rag_client.py

@@ -1,13 +1,13 @@
-#!/usr/bin/env python3
 
 import pulsar
 import _pulsar
 from pulsar.schema import JsonSchema
-from . schema import GraphRagQuery, GraphRagResponse
-from . schema import graph_rag_request_queue, graph_rag_response_queue
+from .. schema import GraphRagQuery, GraphRagResponse
+from .. schema import graph_rag_request_queue, graph_rag_response_queue
 
 import hashlib
 import uuid
+import time
 
 # Ugly
 ERROR=_pulsar.LoggerLevel.Error
@@ -50,9 +50,14 @@ class GraphRagClient:
         )
         self.producer.send(r, properties={ "id": id })
 
-        while True:
+        end_time = time.time() + timeout
 
-            msg = self.consumer.receive(timeout_millis=timeout * 1000)
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
 
             mid = msg.properties()["id"]
 
@@ -64,6 +69,8 @@ class GraphRagClient:
             # Ignore messages with wrong ID
             self.consumer.acknowledge(msg)
 
+        raise TimeoutError("Timed out waiting for response")
+
     def __del__(self):
 
         if hasattr(self, "consumer"):

trustgraph/clients/llm_client.py

@@ -1,14 +1,14 @@
-#!/usr/bin/env python3
 
 import pulsar
 import _pulsar
 from pulsar.schema import JsonSchema
 import hashlib
 import uuid
+import time
 
-from . schema import TextCompletionRequest, TextCompletionResponse
-from . schema import text_completion_request_queue
-from . schema import text_completion_response_queue
+from .. schema import TextCompletionRequest, TextCompletionResponse
+from .. schema import text_completion_request_queue
+from .. schema import text_completion_response_queue
 
 # Ugly
 ERROR=_pulsar.LoggerLevel.Error
@@ -59,11 +59,16 @@ class LlmClient:
             prompt=prompt
         )
 
+        end_time = time.time() + timeout
+
         self.producer.send(r, properties={ "id": id })
 
-        while True:
+        while time.time() < end_time:
 
-            msg = self.consumer.receive(timeout_millis=timeout * 1000)
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
 
             mid = msg.properties()["id"]
 
@@ -75,6 +80,8 @@ class LlmClient:
             # Ignore messages with wrong ID
             self.consumer.acknowledge(msg)
 
+        raise TimeoutError("Timed out waiting for response")
+
     def __del__(self):
 
         if hasattr(self, "consumer"):

trustgraph/clients/prompt_client.py

@@ -1,14 +1,14 @@
-#!/usr/bin/env python3
 
 import pulsar
 import _pulsar
 from pulsar.schema import JsonSchema
 import hashlib
 import uuid
+import time
 
-from . schema import PromptRequest, PromptResponse, Fact
-from . schema import prompt_request_queue
-from . schema import prompt_response_queue
+from .. schema import PromptRequest, PromptResponse, Fact
+from .. schema import prompt_request_queue
+from .. schema import prompt_response_queue
 
 # Ugly
 ERROR=_pulsar.LoggerLevel.Error
@@ -62,9 +62,14 @@ class PromptClient:
 
         self.producer.send(r, properties={ "id": id })
 
-        while True:
+        end_time = time.time() + timeout
 
-            msg = self.consumer.receive(timeout_millis=timeout * 1000)
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
 
             mid = msg.properties()["id"]
 
@@ -76,6 +81,8 @@ class PromptClient:
             # Ignore messages with wrong ID
             self.consumer.acknowledge(msg)
 
+        raise TimeoutError("Timed out waiting for response")
+
     def request_relationships(self, chunk, timeout=500):
 
         id = str(uuid.uuid4())

trustgraph/clients/triples_query_client.py

@@ -5,10 +5,11 @@ import _pulsar
 from pulsar.schema import JsonSchema
 import hashlib
 import uuid
+import time
 
-from . schema import TriplesQueryRequest, TriplesQueryResponse, Value
-from . schema import triples_request_queue
-from . schema import triples_response_queue
+from .. schema import TriplesQueryRequest, TriplesQueryResponse, Value
+from .. schema import triples_request_queue
+from .. schema import triples_response_queue
 
 # Ugly
 ERROR=_pulsar.LoggerLevel.Error
@@ -73,9 +74,14 @@ class TriplesQueryClient:
 
         self.producer.send(r, properties={ "id": id })
 
-        while True:
+        end_time = time.time() + timeout
 
-            msg = self.consumer.receive(timeout_millis=timeout * 1000)
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
 
             mid = msg.properties()["id"]
 
@@ -87,6 +93,8 @@ class TriplesQueryClient:
             # Ignore messages with wrong ID
             self.consumer.acknowledge(msg)
 
+        raise TimeoutError("Timed out waiting for response")
+
     def __del__(self):
 
         if hasattr(self, "consumer"):

trustgraph/embeddings/vectorize/vectorize.py

@@ -7,7 +7,7 @@ Input is text chunk, output is chunk and vectors.
 from ... schema import Chunk, ChunkEmbeddings
 from ... schema import chunk_ingest_queue, chunk_embeddings_ingest_queue
 from ... schema import embeddings_request_queue, embeddings_response_queue
-from ... embeddings_client import EmbeddingsClient
+from ... clients.embeddings_client import EmbeddingsClient
 from ... log_level import LogLevel
 from ... base import ConsumerProducer
 

trustgraph/graph_rag.py

@@ -1,8 +1,8 @@
 
-from . graph_embeddings_client import GraphEmbeddingsClient
-from . triples_query_client import TriplesQueryClient
-from . embeddings_client import EmbeddingsClient
-from . prompt_client import PromptClient
+from . clients.graph_embeddings_client import GraphEmbeddingsClient
+from . clients.triples_query_client import TriplesQueryClient
+from . clients.embeddings_client import EmbeddingsClient
+from . clients.prompt_client import PromptClient
 
 from . schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
 from . schema import TriplesQueryRequest, TriplesQueryResponse

trustgraph/kg/extract_definitions/extract.py

@@ -12,7 +12,7 @@ from ... schema import chunk_embeddings_ingest_queue, triples_store_queue
 from ... schema import prompt_request_queue
 from ... schema import prompt_response_queue
 from ... log_level import LogLevel
-from ... prompt_client import PromptClient
+from ... clients.prompt_client import PromptClient
 from ... rdf import TRUSTGRAPH_ENTITIES, DEFINITION
 from ... base import ConsumerProducer
 

trustgraph/kg/extract_relationships/extract.py

@@ -15,7 +15,7 @@ from ... schema import graph_embeddings_store_queue
 from ... schema import prompt_request_queue
 from ... schema import prompt_response_queue
 from ... log_level import LogLevel
-from ... prompt_client import PromptClient
+from ... clients.prompt_client import PromptClient
 from ... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES
 from ... base import ConsumerProducer
 

trustgraph/model/prompt/generic/service.py

@@ -12,7 +12,7 @@ from .... schema import text_completion_request_queue
 from .... schema import text_completion_response_queue
 from .... schema import prompt_request_queue, prompt_response_queue
 from .... base import ConsumerProducer
-from .... llm_client import LlmClient
+from .... clients.llm_client import LlmClient
 
 from . prompts import to_definitions, to_relationships, to_kg_query