Basic metering module structure

prometheus/prometheus.yml

@@ -29,6 +29,7 @@ scrape_configs:
         - 'kg-extract-definitions:8000'
         - 'kg-extract-topics:8000'
         - 'kg-extract-relationships:8000'
+        - 'metering:8000'
         - 'store-graph-embeddings:8000'
         - 'store-triples:8000'
         - 'text-completion:8000'

scripts/metering

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+from trustgraph.metering import run
+
+run()

setup.py

@@ -81,6 +81,7 @@ setuptools.setup(
         "scripts/load-pdf",
         "scripts/load-text",
         "scripts/load-triples",
+        "scripts/metering",
         "scripts/object-extract-row",
         "scripts/oe-write-milvus",
         "scripts/pdf-decoder",

trustgraph/metering/__init__.py

@@ -0,0 +1,3 @@
+
+from . counter import *
+

trustgraph/metering/__main__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from . counter import run
+
+if __name__ == '__main__':
+    run()
+

trustgraph/metering/counter.py

@@ -0,0 +1,57 @@
+"""
+Simple token counter for each LLM response.
+"""
+
+from prometheus_client import Histogram, Info
+
+from .. schema import TextCompletionResponse, Error
+from .. schema import text_completion_response_queue
+from .. log_level import LogLevel
+from .. base import Consumer
+
+module = ".".join(__name__.split(".")[1:-1])
+
+default_input_queue = text_completion_response_queue
+default_subscriber = module
+
+
+class Processor(Consumer):
+
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+
+        super(Processor, self).__init__(
+            **params | {
+                "input_queue": input_queue,
+                "subscriber": subscriber,
+                "input_schema": TextCompletionResponse,
+            }
+        )
+
+    def handle(self, msg):
+
+        v = msg.value()
+
+        # Sender-produced ID
+        id = msg.properties()["id"]
+
+        print(f"Handling response {id}...", flush=True)
+
+        num_in = v.in_token
+        num_out = v.out_token
+
+        print(f"Input Tokens: {num_in}")
+        print(f"Output Tokens: {num_out}")
+
+    @staticmethod
+    def add_args(parser):
+
+        Consumer.add_args(
+            parser, default_input_queue, default_subscriber,
+        )
+
+def run():
+
+    Processor.start(module, __doc__)

trustgraph/schema/models.py

@@ -1,5 +1,5 @@
 
-from pulsar.schema import Record, String, Array, Double
+from pulsar.schema import Record, String, Array, Double, Integer
 
 from . topic import topic
 from . types import Error
@@ -14,6 +14,8 @@ class TextCompletionRequest(Record):
 class TextCompletionResponse(Record):
     error = Error()
     response = String()
+    in_token = Integer()
+    out_token = Integer()
 
 text_completion_request_queue = topic(
     'text-completion', kind='non-persistent', namespace='request'