Basic metrics working

setup.py

@@ -43,6 +43,7 @@ setuptools.setup(
         "anthropic",
         "google-cloud-aiplatform",
         "pyyaml",
+        "prometheus-client",
     ],
     scripts=[
         "scripts/chunker-recursive",

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__(
+    def __init__(self, **params):
-            self,
-            pulsar_host=default_pulsar_host,
-            log_level=LogLevel.INFO,
-    ):
 
         self.client = None
 
-        if pulsar_host == None:
+        if not hasattr(__class__, "params_metric"):
-            pulsar_host = default_pulsar_host
+            __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,19 +119,13 @@ class BaseProcessor:
 
 class Consumer(BaseProcessor):
 
-    def __init__(
+    def __init__(self, **params):
-            self,
-            pulsar_host=None,
-            log_level=LogLevel.INFO,
-            input_queue="input",
-            subscriber="subscriber",
-            input_schema=None,
-    ):
 
-        super(Consumer, self).__init__(
+        super(Consumer, self).__init__(**params)
-            pulsar_host=pulsar_host,
+
-            log_level=log_level,
+        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")
@@ -144,21 +174,38 @@ class Consumer(BaseProcessor):
 
 class ConsumerProducer(BaseProcessor):
 
-    def __init__(
+    def __init__(self, **params):
-            self,
-            pulsar_host=None,
-            log_level=LogLevel.INFO,
-            input_queue="input",
-            output_queue="output",
-            subscriber="subscriber",
-            input_schema=None,
-            output_schema=None,
-    ):
 
-        super(ConsumerProducer, self).__init__(
+        input_queue = params.get("input_queue")
-            pulsar_host=pulsar_host,
+        output_queue = params.get("output_queue")
-            log_level=log_level,
+        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__, "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 +231,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,6 +246,8 @@ class ConsumerProducer(BaseProcessor):
                 # Message failed to be processed
                 self.consumer.negative_acknowledge(msg)
 
+                __class__.processing_metric.labels(status="error").inc()
+
     def send(self, msg, properties={}):
 
         self.producer.send(msg, properties)
@@ -228,18 +280,12 @@ class ConsumerProducer(BaseProcessor):
 
 class Producer(BaseProcessor):
 
-    def __init__(
+    def __init__(self, **params):
-            self,
-            pulsar_host=None,
-            log_level=LogLevel.INFO,
-            output_queue="output",
-            output_schema=None,
-    ):
 
-        super(Producer, self).__init__(
+        output_queue = params.get("output_queue")
-            pulsar_host=pulsar_host,
+        output_schema = params.get("output_schema")
-            log_level=log_level,
+
-        )
+        super(Producer, self).__init__(**params)
 
         if output_schema == None:
             raise RuntimeError("output_schema must be specified")

trustgraph/llm/ollama_text/llm.py

@@ -5,6 +5,7 @@ Input is prompt, output is response.
 """
 
 from langchain_community.llms import Ollama
+from prometheus_client import Histogram, Info, Counter
 
 from ... schema import TextCompletionRequest, TextCompletionResponse
 from ... log_level import LogLevel
@@ -18,27 +19,36 @@ default_ollama = 'http://localhost:11434'
 
 class Processor(ConsumerProducer):
 
-    def __init__(
+    def __init__(self, **params):
-            self,
+
-            pulsar_host=None,
+        input_queue = params.get("input_queue", default_input_queue)
-            input_queue=default_input_queue,
+        output_queue = params.get("output_queue", default_output_queue)
-            output_queue=default_output_queue,
+        subscriber = params.get("subscriber", default_subscriber)
-            subscriber=default_subscriber,
+        model = params.get("model", default_model)
-            log_level=LogLevel.INFO,
+        ollama = params.get("ollama", default_ollama)
-            model=default_model,
-            ollama=default_ollama,
-    ):
 
         super(Processor, self).__init__(
-            pulsar_host=pulsar_host,
+            **params | {
-            log_level=log_level,
+                "input_queue": input_queue,
-            input_queue=input_queue,
+                "output_queue": output_queue,
-            output_queue=output_queue,
+                "subscriber": subscriber,
-            subscriber=subscriber,
+                "model": model,
-            input_schema=TextCompletionRequest,
+                "ollama": ollama,
-            output_schema=TextCompletionResponse,
+                "input_schema": TextCompletionRequest,
+                "output_schema": TextCompletionResponse,
+            }
         )
 
+        if not hasattr(__class__, "model_metric"):
+            __class__.model_metric = Info(
+                'model', 'Model information'
+            )
+
+        __class__.model_metric.info({
+            "model": model,
+            "ollama": ollama,
+        })
+
         self.llm = Ollama(base_url=ollama, model=model)
 
     def handle(self, msg):