trustgraph 0.3.1__tar.gz → 0.4.2__tar.gz

{trustgraph-0.3.1 → trustgraph-0.4.2}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph
-Version: 0.3.1
+Version: 0.4.2
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Home-page: https://github.com/trustgraph-ai/trustgraph
-Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.3.1.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.4.2.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3
@@ -31,6 +31,7 @@ Requires-Dist: pypdf
 Requires-Dist: anthropic
 Requires-Dist: google-cloud-aiplatform
 Requires-Dist: pyyaml
+Requires-Dist: prometheus-client
 
 
 # TrustGraph

{trustgraph-0.3.1 → trustgraph-0.4.2}/setup.py

@@ -4,7 +4,7 @@ import os
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
-version = "0.3.1"
+version = "0.4.2"
 
 setuptools.setup(
     name="trustgraph",
@@ -43,6 +43,7 @@ setuptools.setup(
         "anthropic",
         "google-cloud-aiplatform",
         "pyyaml",
+        "prometheus-client",
     ],
     scripts=[
         "scripts/chunker-recursive",

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/base/processor.py

@@ -2,8 +2,10 @@
 import os
 import argparse
 import pulsar
+import _pulsar
 import time
 from pulsar.schema import JsonSchema
+from prometheus_client import start_http_server, Histogram, Info, Counter
 
 from .. log_level import LogLevel
 
@@ -11,16 +13,23 @@ class BaseProcessor:
 
     default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
 
-    def __init__(
-            self,
-            pulsar_host=default_pulsar_host,
-            log_level=LogLevel.INFO,
-    ):
+    def __init__(self, **params):
 
         self.client = None
 
-        if pulsar_host == None:
-            pulsar_host = default_pulsar_host
+        if not hasattr(__class__, "params_metric"):
+            __class__.params_metric = Info(
+                'params', 'Parameters configuration'
+            )
+
+        # FIXME: Maybe outputs information it should not
+        __class__.params_metric.info({
+            k: str(params[k])
+            for k in params
+        })
+
+        pulsar_host = params.get("pulsar_host", self.default_pulsar_host)
+        log_level = params.get("log_level", LogLevel.INFO)
 
         self.pulsar_host = pulsar_host
 
@@ -51,6 +60,20 @@ class BaseProcessor:
             help=f'Output queue (default: info)'
         )
 
+        parser.add_argument(
+            '-M', '--metrics-enabled',
+            type=bool,
+            default=True,
+            help=f'Pulsar host (default: true)',
+        )
+
+        parser.add_argument(
+            '-P', '--metrics-port',
+            type=int,
+            default=8000,
+            help=f'Pulsar host (default: 8000)',
+        )
+
     def run(self):
         raise RuntimeError("Something should have implemented the run method")
 
@@ -69,13 +92,26 @@ class BaseProcessor:
             args = parser.parse_args()
             args = vars(args)
 
+            if args["metrics_enabled"]:
+                start_http_server(args["metrics_port"])
+
             try:
 
                 p = cls(**args)
                 p.run()
 
+            except KeyboardInterrupt:
+                print("Keyboard interrupt.")
+                return
+
+            except _pulsar.Interrupted:
+                print("Pulsar Interrupted.")
+                return
+
             except Exception as e:
 
+                print(type(e))
+
                 print("Exception:", e, flush=True)
                 print("Will retry...", flush=True)
 
@@ -83,23 +119,38 @@ class BaseProcessor:
 
 class Consumer(BaseProcessor):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            log_level=LogLevel.INFO,
-            input_queue="input",
-            subscriber="subscriber",
-            input_schema=None,
-    ):
+    def __init__(self, **params):
 
-        super(Consumer, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-        )
+        super(Consumer, self).__init__(**params)
+
+        input_queue = params.get("input_queue")
+        subscriber = params.get("subscriber")
+        input_schema = params.get("input_schema")
 
         if input_schema == None:
             raise RuntimeError("input_schema must be specified")
 
+        if not hasattr(__class__, "request_metric"):
+            __class__.request_metric = Histogram(
+                'request_latency', 'Request latency (seconds)'
+            )
+
+        if not hasattr(__class__, "pubsub_metric"):
+            __class__.pubsub_metric = Info(
+                'pubsub', 'Pub/sub configuration'
+            )
+
+        if not hasattr(__class__, "processing_metric"):
+            __class__.processing_metric = Counter(
+                'processing_count', 'Processing count', ["status"]
+            )
+
+        __class__.pubsub_metric.info({
+            "input_queue": input_queue,
+            "subscriber": subscriber,
+            "input_schema": input_schema.__name__,
+        })
+
         self.consumer = self.client.subscribe(
             input_queue, subscriber,
             schema=JsonSchema(input_schema),
@@ -113,11 +164,14 @@ class Consumer(BaseProcessor):
 
             try:
 
-                self.handle(msg)
+                with __class__.request_metric.time():
+                    self.handle(msg)
 
                 # Acknowledge successful processing of the message
                 self.consumer.acknowledge(msg)
 
+                __class__.processing_metric.labels(status="success").inc()
+
             except Exception as e:
 
                 print("Exception:", e, flush=True)
@@ -125,6 +179,8 @@ class Consumer(BaseProcessor):
                 # Message failed to be processed
                 self.consumer.negative_acknowledge(msg)
 
+                __class__.processing_metric.labels(status="error").inc()
+
     @staticmethod
     def add_args(parser, default_input_queue, default_subscriber):
 
@@ -144,21 +200,43 @@ class Consumer(BaseProcessor):
 
 class ConsumerProducer(BaseProcessor):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            log_level=LogLevel.INFO,
-            input_queue="input",
-            output_queue="output",
-            subscriber="subscriber",
-            input_schema=None,
-            output_schema=None,
-    ):
+    def __init__(self, **params):
 
-        super(ConsumerProducer, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-        )
+        input_queue = params.get("input_queue")
+        output_queue = params.get("output_queue")
+        subscriber = params.get("subscriber")
+        input_schema = params.get("input_schema")
+        output_schema = params.get("output_schema")
+
+        if not hasattr(__class__, "request_metric"):
+            __class__.request_metric = Histogram(
+                'request_latency', 'Request latency (seconds)'
+            )
+
+        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'
+            )
+
+        if not hasattr(__class__, "processing_metric"):
+            __class__.processing_metric = Counter(
+                'processing_count', 'Processing count', ["status"]
+            )
+
+        __class__.pubsub_metric.info({
+            "input_queue": input_queue,
+            "output_queue": output_queue,
+            "subscriber": subscriber,
+            "input_schema": input_schema.__name__,
+            "output_schema": output_schema.__name__,
+        })
+
+        super(ConsumerProducer, self).__init__(**params)
 
         if input_schema == None:
             raise RuntimeError("input_schema must be specified")
@@ -184,11 +262,14 @@ class ConsumerProducer(BaseProcessor):
 
             try:
 
-                resp = self.handle(msg)
+                with __class__.request_metric.time():
+                    resp = self.handle(msg)
 
                 # Acknowledge successful processing of the message
                 self.consumer.acknowledge(msg)
 
+                __class__.processing_metric.labels(status="success").inc()
+
             except Exception as e:
 
                 print("Exception:", e, flush=True)
@@ -196,9 +277,11 @@ class ConsumerProducer(BaseProcessor):
                 # Message failed to be processed
                 self.consumer.negative_acknowledge(msg)
 
-    def send(self, msg, properties={}):
+                __class__.processing_metric.labels(status="error").inc()
 
+    def send(self, msg, properties={}):
         self.producer.send(msg, properties)
+        __class__.output_metric.inc()
 
     @staticmethod
     def add_args(
@@ -228,18 +311,27 @@ class ConsumerProducer(BaseProcessor):
 
 class Producer(BaseProcessor):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            log_level=LogLevel.INFO,
-            output_queue="output",
-            output_schema=None,
-    ):
+    def __init__(self, **params):
 
-        super(Producer, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-        )
+        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")
@@ -250,8 +342,8 @@ class Producer(BaseProcessor):
         )
 
     def send(self, msg, properties={}):
-
         self.producer.send(msg, properties)
+        __class__.output_metric.inc()
 
     @staticmethod
     def add_args(

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/chunker/recursive/chunker.py

@@ -17,25 +17,22 @@ default_subscriber = 'chunker-recursive'
 
 class Processor(ConsumerProducer):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            input_queue=default_input_queue,
-            output_queue=default_output_queue,
-            subscriber=default_subscriber,
-            log_level=LogLevel.INFO,
-            chunk_size=2000,
-            chunk_overlap=100,
-    ):
-
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+        chunk_size = params.get("chunk_size", 2000)
+        chunk_overlap = params.get("chunk_overlap", 100)
+        
         super(Processor, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-            input_queue=input_queue,
-            output_queue=output_queue,
-            subscriber=subscriber,
-            input_schema=TextDocument,
-            output_schema=Chunk,
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": TextDocument,
+                "output_schema": Chunk,
+            }
         )
 
         self.text_splitter = RecursiveCharacterTextSplitter(

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/decoder/pdf/pdf_decoder.py

@@ -18,23 +18,20 @@ default_subscriber = 'pdf-decoder'
 
 class Processor(ConsumerProducer):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            input_queue=default_input_queue,
-            output_queue=default_output_queue,
-            subscriber=default_subscriber,
-            log_level=LogLevel.INFO,
-    ):
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
 
         super(Processor, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-            input_queue=input_queue,
-            output_queue=output_queue,
-            subscriber=subscriber,
-            input_schema=Document,
-            output_schema=TextDocument,
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": Document,
+                "output_schema": TextDocument,
+            }
         )
 
         print("PDF inited")

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/embeddings/hf/hf.py

@@ -17,24 +17,21 @@ default_model="all-MiniLM-L6-v2"
 
 class Processor(ConsumerProducer):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            input_queue=default_input_queue,
-            output_queue=default_output_queue,
-            subscriber=default_subscriber,
-            log_level=LogLevel.INFO,
-            model=default_model,
-    ):
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+        model = params.get("model", default_model)
 
         super(Processor, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-            input_queue=input_queue,
-            output_queue=output_queue,
-            subscriber=subscriber,
-            input_schema=EmbeddingsRequest,
-            output_schema=EmbeddingsResponse,
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": EmbeddingsRequest,
+                "output_schema": EmbeddingsResponse,
+            }
         )
 
         self.embeddings = HuggingFaceEmbeddings(model_name=model)

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/embeddings/ollama/processor.py

@@ -17,25 +17,20 @@ default_ollama = 'http://localhost:11434'
 
 class Processor(ConsumerProducer):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            input_queue=default_input_queue,
-            output_queue=default_output_queue,
-            subscriber=default_subscriber,
-            log_level=LogLevel.INFO,
-            model=default_model,
-            ollama=default_ollama,
-    ):
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
 
         super(Processor, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-            input_queue=input_queue,
-            output_queue=output_queue,
-            subscriber=subscriber,
-            input_schema=EmbeddingsRequest,
-            output_schema=EmbeddingsResponse,
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": EmbeddingsRequest,
+                "output_schema": EmbeddingsResponse,
+            }
         )
 
         self.embeddings = OllamaEmbeddings(base_url=ollama, model=model)

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/embeddings/vectorize/vectorize.py

@@ -15,26 +15,23 @@ default_subscriber = 'embeddings-vectorizer'
 
 class Processor(ConsumerProducer):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            input_queue=default_input_queue,
-            output_queue=default_output_queue,
-            subscriber=default_subscriber,
-            log_level=LogLevel.INFO,
-    ):
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
 
         super(Processor, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-            input_queue=input_queue,
-            output_queue=output_queue,
-            subscriber=subscriber,
-            input_schema=Chunk,
-            output_schema=VectorsChunk,
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": Chunk,
+                "output_schema": VectorsChunk,
+            }
         )
 
-        self.embeddings = EmbeddingsClient(pulsar_host=pulsar_host)
+        self.embeddings = EmbeddingsClient(pulsar_host=self.pulsar_host)
 
     def emit(self, source, chunk, vectors):
 

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/graph/cassandra_write/write.py

@@ -20,27 +20,22 @@ default_graph_host='localhost'
 
 class Processor(Consumer):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            input_queue=default_input_queue,
-            subscriber=default_subscriber,
-            graph_host=default_graph_host,
-            log_level=LogLevel.INFO,
-    ):
+    def __init__(self, **params):
+        
+        input_queue = params.get("input_queue", default_input_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+        graph_host = params.get("graph_host", default_graph_host)
 
         super(Processor, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-            input_queue=input_queue,
-            subscriber=subscriber,
-            input_schema=Triple,
+            **params | {
+                "input_queue": input_queue,
+                "subscriber": subscriber,
+                "input_schema": Triple,
+            }
         )
 
         self.tg = TrustGraph([graph_host])
 
-        self.count = 0
-
     def handle(self, msg):
 
         v = msg.value()
@@ -51,11 +46,6 @@ class Processor(Consumer):
             v.o.value
         )
 
-        self.count += 1
-
-        if (self.count % 1000) == 0:
-            print(self.count, "...", flush=True)
-
     @staticmethod
     def add_args(parser):
 

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/graph_rag.py

@@ -18,7 +18,7 @@ class GraphRag:
             verbose=False,
             entity_limit=50,
             triple_limit=30,
-            max_sg_size=3000,
+            max_subgraph_size=3000,
     ):
 
         self.verbose=verbose
@@ -37,7 +37,7 @@ class GraphRag:
 
         self.entity_limit=entity_limit
         self.query_limit=triple_limit
-        self.max_sg_size=max_sg_size
+        self.max_subgraph_size=max_subgraph_size
 
         self.label_cache = {}
 
@@ -149,7 +149,7 @@ class GraphRag:
 
         subgraph = list(subgraph)
 
-        subgraph = subgraph[0:self.max_sg_size]
+        subgraph = subgraph[0:self.max_subgraph_size]
 
         if self.verbose:
             print("Subgraph:", flush=True)

{trustgraph-0.3.1 → trustgraph-0.4.2}/trustgraph/kg/extract_definitions/extract.py

@@ -22,26 +22,23 @@ default_subscriber = 'kg-extract-definitions'
 
 class Processor(ConsumerProducer):
 
-    def __init__(
-            self,
-            pulsar_host=None,
-            input_queue=default_input_queue,
-            output_queue=default_output_queue,
-            subscriber=default_subscriber,
-            log_level=LogLevel.INFO,
-    ):
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
 
         super(Processor, self).__init__(
-            pulsar_host=pulsar_host,
-            log_level=log_level,
-            input_queue=input_queue,
-            output_queue=output_queue,
-            subscriber=subscriber,
-            input_schema=VectorsChunk,
-            output_schema=Triple,
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": VectorsChunk,
+                "output_schema": Triple,
+            }
         )
 
-        self.llm = LlmClient(pulsar_host=pulsar_host)
+        self.llm = LlmClient(pulsar_host=self.pulsar_host)
 
     def to_uri(self, text):