From 74fcba1b3e28bed1250568e31117bb822cf19dfd Mon Sep 17 00:00:00 2001
From: JackColquitt <daniel@kalntera.ai>
Date: Fri, 27 Sep 2024 16:57:46 -0700
Subject: [PATCH] Basic metering module structure

---
 prometheus/prometheus.yml       |  1 +
 scripts/metering                |  5 +++
 setup.py                        |  1 +
 trustgraph/metering/__init__.py |  3 ++
 trustgraph/metering/__main__.py |  7 ++++
 trustgraph/metering/counter.py  | 57 +++++++++++++++++++++++++++++++++
 trustgraph/schema/models.py     |  4 ++-
 7 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 scripts/metering
 create mode 100644 trustgraph/metering/__init__.py
 create mode 100755 trustgraph/metering/__main__.py
 create mode 100644 trustgraph/metering/counter.py

diff --git a/prometheus/prometheus.yml b/prometheus/prometheus.yml
index f459dad9..60a8e488 100644
--- a/prometheus/prometheus.yml
+++ b/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'
diff --git a/scripts/metering b/scripts/metering
new file mode 100644
index 00000000..7f1d0e12
--- /dev/null
+++ b/scripts/metering
@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+from trustgraph.metering import run
+
+run()
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 8240a805..e50b4e51 100644
--- a/setup.py
+++ b/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",
diff --git a/trustgraph/metering/__init__.py b/trustgraph/metering/__init__.py
new file mode 100644
index 00000000..0ed03774
--- /dev/null
+++ b/trustgraph/metering/__init__.py
@@ -0,0 +1,3 @@
+
+from . counter import *
+
diff --git a/trustgraph/metering/__main__.py b/trustgraph/metering/__main__.py
new file mode 100755
index 00000000..802f2b8d
--- /dev/null
+++ b/trustgraph/metering/__main__.py
@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from . counter import run
+
+if __name__ == '__main__':
+    run()
+
diff --git a/trustgraph/metering/counter.py b/trustgraph/metering/counter.py
new file mode 100644
index 00000000..82e652d5
--- /dev/null
+++ b/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__)
\ No newline at end of file
diff --git a/trustgraph/schema/models.py b/trustgraph/schema/models.py
index cf73a203..192c51b6 100644
--- a/trustgraph/schema/models.py
+++ b/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'