- Change flawed _client timeout logic which was causing major lags

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

trustgraph/clients/embeddings_client.py

@@ -0,0 +1,95 @@
+
+import pulsar
+import _pulsar
+from pulsar.schema import JsonSchema
+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
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class EmbeddingsClient:
+
+    def __init__(
+            self, log_level=ERROR,
+            input_queue=None,
+            output_queue=None,
+            subscriber=None,
+            pulsar_host="pulsar://pulsar:6650",
+    ):
+
+        self.client = None
+
+        if input_queue == None:
+            input_queue=embeddings_request_queue
+
+        if output_queue == None:
+            output_queue=embeddings_response_queue
+
+        if subscriber == None:
+            subscriber = str(uuid.uuid4())
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level),
+        )
+
+        self.producer = self.client.create_producer(
+            topic=input_queue,
+            schema=JsonSchema(EmbeddingsRequest),
+            chunking_enabled=True,
+        )
+
+        self.consumer = self.client.subscribe(
+            output_queue, subscriber,
+            schema=JsonSchema(EmbeddingsResponse),
+        )
+
+    def request(self, text, timeout=10):
+
+        id = str(uuid.uuid4())
+
+        r = EmbeddingsRequest(
+            text=text
+        )
+        self.producer.send(r, properties={ "id": id })
+
+        end_time = time.time() + timeout
+
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                resp = msg.value().vectors
+                self.consumer.acknowledge(msg)
+                return resp
+
+            # Ignore messages with wrong ID
+            self.consumer.acknowledge(msg)
+
+        raise TimeoutError("Timed out waiting for response")
+
+    def __del__(self):
+
+        if hasattr(self, "consumer"):
+#             self.consumer.unsubscribe()
+            self.consumer.close()
+            
+        if hasattr(self, "producer"):
+            self.producer.flush()
+            self.producer.close()
+            
+        self.client.close()
+

trustgraph/clients/graph_embeddings_client.py

@@ -0,0 +1,96 @@
+
+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
+
+# Ugly
+ERROR=_pulsar.LoggerLevel.Error
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class GraphEmbeddingsClient:
+
+    def __init__(
+            self, log_level=ERROR,
+            subscriber=None,
+            input_queue=None,
+            output_queue=None,
+            pulsar_host="pulsar://pulsar:6650",
+    ):
+
+        if input_queue == None:
+            input_queue = graph_embeddings_request_queue
+
+        if output_queue == None:
+            output_queue = graph_embeddings_response_queue
+
+        if subscriber == None:
+            subscriber = str(uuid.uuid4())
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level),
+        )
+
+        self.producer = self.client.create_producer(
+            topic=input_queue,
+            schema=JsonSchema(GraphEmbeddingsRequest),
+            chunking_enabled=True,
+        )
+
+        self.consumer = self.client.subscribe(
+            output_queue, subscriber,
+            schema=JsonSchema(GraphEmbeddingsResponse),
+        )
+
+    def request(self, vectors, limit=10, timeout=500):
+
+        id = str(uuid.uuid4())
+
+        r = GraphEmbeddingsRequest(
+            vectors=vectors,
+            limit=limit,
+        )
+
+        self.producer.send(r, properties={ "id": id })
+
+        end_time = time.time() + timeout
+
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                resp = msg.value().entities
+                self.consumer.acknowledge(msg)
+                return resp
+
+            # Ignore messages with wrong ID
+            self.consumer.acknowledge(msg)
+
+        raise TimeoutError("Timed out waiting for response")
+
+    def __del__(self):
+
+        if hasattr(self, "consumer"):
+            self.consumer.close()
+            
+        if hasattr(self, "producer"):
+            self.producer.flush()
+            self.producer.close()
+            
+        self.client.close()
+

trustgraph/clients/graph_rag_client.py

@@ -0,0 +1,85 @@
+
+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
+
+import hashlib
+import uuid
+import time
+
+# Ugly
+ERROR=_pulsar.LoggerLevel.Error
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class GraphRagClient:
+
+    def __init__(
+            self, log_level=ERROR, subscriber=None,
+            pulsar_host="pulsar://pulsar:6650",
+    ):
+
+        if subscriber == None:
+            subscriber = str(uuid.uuid4())
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level),
+        )
+
+        self.producer = self.client.create_producer(
+            topic=graph_rag_request_queue,
+            schema=JsonSchema(GraphRagQuery),
+            chunking_enabled=True,
+        )
+
+        self.consumer = self.client.subscribe(
+            graph_rag_response_queue, subscriber,
+            schema=JsonSchema(GraphRagResponse),
+        )
+
+    def request(self, query, timeout=500):
+
+        id = str(uuid.uuid4())
+
+        r = GraphRagQuery(
+            query=query
+        )
+        self.producer.send(r, properties={ "id": id })
+
+        end_time = time.time() + timeout
+
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                resp = msg.value().response
+                self.consumer.acknowledge(msg)
+                return resp
+
+            # Ignore messages with wrong ID
+            self.consumer.acknowledge(msg)
+
+        raise TimeoutError("Timed out waiting for response")
+
+    def __del__(self):
+
+        if hasattr(self, "consumer"):
+#             self.consumer.unsubscribe()
+            self.consumer.close()
+            
+        if hasattr(self, "producer"):
+            self.producer.flush()
+            self.producer.close()
+            
+        self.client.close()
+

trustgraph/clients/llm_client.py

@@ -0,0 +1,96 @@
+
+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
+
+# Ugly
+ERROR=_pulsar.LoggerLevel.Error
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class LlmClient:
+
+    def __init__(
+            self, log_level=ERROR,
+            subscriber=None,
+            input_queue=None,
+            output_queue=None,
+            pulsar_host="pulsar://pulsar:6650",
+    ):
+
+        if input_queue == None:
+            input_queue = text_completion_request_queue
+
+        if output_queue == None:
+            output_queue = text_completion_response_queue
+
+        if subscriber == None:
+            subscriber = str(uuid.uuid4())
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level),
+        )
+
+        self.producer = self.client.create_producer(
+            topic=input_queue,
+            schema=JsonSchema(TextCompletionRequest),
+            chunking_enabled=True,
+        )
+
+        self.consumer = self.client.subscribe(
+            output_queue, subscriber,
+            schema=JsonSchema(TextCompletionResponse),
+        )
+
+    def request(self, prompt, timeout=30):
+
+        id = str(uuid.uuid4())
+
+        r = TextCompletionRequest(
+            prompt=prompt
+        )
+
+        end_time = time.time() + timeout
+
+        self.producer.send(r, properties={ "id": id })
+
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                resp = msg.value().response
+                self.consumer.acknowledge(msg)
+                return resp
+
+            # Ignore messages with wrong ID
+            self.consumer.acknowledge(msg)
+
+        raise TimeoutError("Timed out waiting for response")
+
+    def __del__(self):
+
+        if hasattr(self, "consumer"):
+#             self.consumer.unsubscribe()
+            self.consumer.close()
+            
+        if hasattr(self, "producer"):
+            self.producer.flush()
+            self.producer.close()
+            
+        self.client.close()
+

trustgraph/clients/prompt_client.py

@@ -0,0 +1,150 @@
+
+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
+
+# Ugly
+ERROR=_pulsar.LoggerLevel.Error
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class PromptClient:
+
+    def __init__(
+            self, log_level=ERROR,
+            subscriber=None,
+            input_queue=None,
+            output_queue=None,
+            pulsar_host="pulsar://pulsar:6650",
+    ):
+
+        if input_queue == None:
+            input_queue = prompt_request_queue
+
+        if output_queue == None:
+            output_queue = prompt_response_queue
+
+        if subscriber == None:
+            subscriber = str(uuid.uuid4())
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level),
+        )
+
+        self.producer = self.client.create_producer(
+            topic=input_queue,
+            schema=JsonSchema(PromptRequest),
+            chunking_enabled=True,
+        )
+
+        self.consumer = self.client.subscribe(
+            output_queue, subscriber,
+            schema=JsonSchema(PromptResponse),
+        )
+
+    def request_definitions(self, chunk, timeout=30):
+
+        id = str(uuid.uuid4())
+
+        r = PromptRequest(
+            kind="extract-definitions",
+            chunk=chunk,
+        )
+
+        self.producer.send(r, properties={ "id": id })
+
+        end_time = time.time() + timeout
+
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                resp = msg.value().definitions
+                self.consumer.acknowledge(msg)
+                return resp
+
+            # 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())
+
+        r = PromptRequest(
+            kind="extract-relationships",
+            chunk=chunk,
+        )
+
+        self.producer.send(r, properties={ "id": id })
+
+        while True:
+
+            msg = self.consumer.receive(timeout_millis=timeout * 1000)
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                resp = msg.value().relationships
+                self.consumer.acknowledge(msg)
+                return resp
+
+            # Ignore messages with wrong ID
+            self.consumer.acknowledge(msg)
+
+    def request_kg_prompt(self, query, kg, timeout=500):
+
+        id = str(uuid.uuid4())
+
+        r = PromptRequest(
+            kind="kg-prompt",
+            query=query,
+            kg=[
+                Fact(s=v[0], p=v[1], o=v[2])
+                for v in kg
+            ],
+        )
+
+        self.producer.send(r, properties={ "id": id })
+
+        while True:
+
+            msg = self.consumer.receive(timeout_millis=timeout * 1000)
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                resp = msg.value().answer
+                self.consumer.acknowledge(msg)
+                return resp
+
+            # Ignore messages with wrong ID
+            self.consumer.acknowledge(msg)
+
+    def __del__(self):
+
+        if hasattr(self, "consumer"):
+            self.consumer.close()
+            
+        if hasattr(self, "producer"):
+            self.producer.flush()
+            self.producer.close()
+            
+        self.client.close()
+

trustgraph/clients/triples_query_client.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+
+import pulsar
+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
+
+# Ugly
+ERROR=_pulsar.LoggerLevel.Error
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class TriplesQueryClient:
+
+    def __init__(
+            self, log_level=ERROR,
+            subscriber=None,
+            input_queue=None,
+            output_queue=None,
+            pulsar_host="pulsar://pulsar:6650",
+    ):
+
+        if input_queue == None:
+            input_queue = triples_request_queue
+
+        if output_queue == None:
+            output_queue = triples_response_queue
+
+        if subscriber == None:
+            subscriber = str(uuid.uuid4())
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level),
+        )
+
+        self.producer = self.client.create_producer(
+            topic=input_queue,
+            schema=JsonSchema(TriplesQueryRequest),
+            chunking_enabled=True,
+        )
+
+        self.consumer = self.client.subscribe(
+            output_queue, subscriber,
+            schema=JsonSchema(TriplesQueryResponse),
+        )
+
+    def create_value(self, ent):
+
+        if ent == None: return None
+
+        if ent.startswith("http://") or ent.startswith("https://"):
+            return Value(value=ent, is_uri=True)
+
+        return Value(value=ent, is_uri=False)
+
+    def request(self, s, p, o, limit=10, timeout=500):
+
+        id = str(uuid.uuid4())
+
+        r = TriplesQueryRequest(
+            s=self.create_value(s),
+            p=self.create_value(p),
+            o=self.create_value(o),
+            limit=limit,
+        )
+
+        self.producer.send(r, properties={ "id": id })
+
+        end_time = time.time() + timeout
+
+        while time.time() < end_time:
+
+            try:
+                msg = self.consumer.receive(timeout_millis=5000)
+            except pulsar.exceptions.Timeout:
+                continue
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                resp = msg.value().triples
+                self.consumer.acknowledge(msg)
+                return resp
+
+            # Ignore messages with wrong ID
+            self.consumer.acknowledge(msg)
+
+        raise TimeoutError("Timed out waiting for response")
+
+    def __del__(self):
+
+        if hasattr(self, "consumer"):
+            self.consumer.close()
+            
+        if hasattr(self, "producer"):
+            self.producer.flush()
+            self.producer.close()
+            
+        self.client.close()
+