From 313767f16cc9178ae4e21191319ab67d300ba60a Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Mon, 30 Sep 2024 23:55:09 +0100
Subject: [PATCH] Added Prom metrics to metering, added dashboard

---
 grafana/dashboards/dashboard.json             | 236 +++++++++++++++++-
 .../trustgraph/metering/counter.py            |  30 ++-
 2 files changed, 262 insertions(+), 4 deletions(-)

diff --git a/grafana/dashboards/dashboard.json b/grafana/dashboards/dashboard.json
index aba5d918..cc979994 100644
--- a/grafana/dashboards/dashboard.json
+++ b/grafana/dashboards/dashboard.json
@@ -824,7 +824,7 @@
       },
       "gridPos": {
         "h": 7,
-        "w": 12,
+        "w": 8,
         "x": 0,
         "y": 32
       },
@@ -895,6 +895,238 @@
         }
       ],
       "type": "table"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "f6b18033-5918-4e05-a1ca-4cb30343b129"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "custom": {
+            "axisBorderShow": false,
+            "axisCenteredZero": false,
+            "axisColorMode": "text",
+            "axisLabel": "",
+            "axisPlacement": "auto",
+            "barAlignment": 0,
+            "drawStyle": "line",
+            "fillOpacity": 0,
+            "gradientMode": "none",
+            "hideFrom": {
+              "legend": false,
+              "tooltip": false,
+              "viz": false
+            },
+            "insertNulls": false,
+            "lineInterpolation": "linear",
+            "lineWidth": 1,
+            "pointSize": 5,
+            "scaleDistribution": {
+              "type": "linear"
+            },
+            "showPoints": "auto",
+            "spanNulls": false,
+            "stacking": {
+              "group": "A",
+              "mode": "none"
+            },
+            "thresholdsStyle": {
+              "mode": "off"
+            }
+          },
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "red",
+                "value": 80
+              }
+            ]
+          }
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 7,
+        "w": 8,
+        "x": 8,
+        "y": 32
+      },
+      "id": 15,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "single",
+          "sort": "none"
+        }
+      },
+      "targets": [
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "f6b18033-5918-4e05-a1ca-4cb30343b129"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "sum by(job) (rate(input_tokens_total[$__rate_interval]))",
+          "fullMetaSearch": false,
+          "includeNullMetadata": true,
+          "instant": false,
+          "legendFormat": "input {{job}}",
+          "range": true,
+          "refId": "A",
+          "useBackend": false
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "f6b18033-5918-4e05-a1ca-4cb30343b129"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "sum by(job) (rate(output_tokens_total[$__rate_interval]))",
+          "fullMetaSearch": false,
+          "hide": false,
+          "includeNullMetadata": true,
+          "instant": false,
+          "legendFormat": "output {{job}}",
+          "range": true,
+          "refId": "B",
+          "useBackend": false
+        }
+      ],
+      "title": "Tokens",
+      "type": "timeseries"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "f6b18033-5918-4e05-a1ca-4cb30343b129"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "custom": {
+            "axisBorderShow": false,
+            "axisCenteredZero": false,
+            "axisColorMode": "text",
+            "axisLabel": "",
+            "axisPlacement": "auto",
+            "barAlignment": 0,
+            "drawStyle": "line",
+            "fillOpacity": 0,
+            "gradientMode": "none",
+            "hideFrom": {
+              "legend": false,
+              "tooltip": false,
+              "viz": false
+            },
+            "insertNulls": false,
+            "lineInterpolation": "linear",
+            "lineWidth": 1,
+            "pointSize": 5,
+            "scaleDistribution": {
+              "type": "linear"
+            },
+            "showPoints": "auto",
+            "spanNulls": false,
+            "stacking": {
+              "group": "A",
+              "mode": "none"
+            },
+            "thresholdsStyle": {
+              "mode": "off"
+            }
+          },
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "red",
+                "value": 80
+              }
+            ]
+          }
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 7,
+        "w": 8,
+        "x": 16,
+        "y": 32
+      },
+      "id": 16,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "single",
+          "sort": "none"
+        }
+      },
+      "targets": [
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "f6b18033-5918-4e05-a1ca-4cb30343b129"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "sum by(job) (rate(input_cost_total[$__rate_interval]))",
+          "fullMetaSearch": false,
+          "includeNullMetadata": true,
+          "instant": false,
+          "legendFormat": "input {{job}}",
+          "range": true,
+          "refId": "A",
+          "useBackend": false
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "f6b18033-5918-4e05-a1ca-4cb30343b129"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "sum by(job) (rate(output_cost_total[$__rate_interval]))",
+          "fullMetaSearch": false,
+          "hide": false,
+          "includeNullMetadata": true,
+          "instant": false,
+          "legendFormat": "output {{job}}",
+          "range": true,
+          "refId": "B",
+          "useBackend": false
+        }
+      ],
+      "title": "Token cost",
+      "type": "timeseries"
     }
   ],
   "refresh": "5s",
@@ -911,6 +1143,6 @@
   "timezone": "",
   "title": "Overview",
   "uid": "b5c8abf8-fe79-496b-b028-10bde917d1f0",
-  "version": 3,
+  "version": 2,
   "weekStart": ""
 }
diff --git a/trustgraph-flow/trustgraph/metering/counter.py b/trustgraph-flow/trustgraph/metering/counter.py
index 6404888e..6e6b829b 100644
--- a/trustgraph-flow/trustgraph/metering/counter.py
+++ b/trustgraph-flow/trustgraph/metering/counter.py
@@ -2,7 +2,7 @@
 Simple token counter for each LLM response.
 """
 
-from prometheus_client import Histogram, Info
+from prometheus_client import Counter
 from . pricelist import price_list
 
 from .. schema import TextCompletionResponse, Error
@@ -20,6 +20,26 @@ class Processor(Consumer):
 
     def __init__(self, **params):
 
+        if not hasattr(__class__, "input_token_metric"):
+            __class__.input_token_metric = Counter(
+                'input_tokens', 'Input token count'
+            )
+
+        if not hasattr(__class__, "output_token_metric"):
+            __class__.output_token_metric = Counter(
+                'output_tokens', 'Output token count'
+            )
+
+        if not hasattr(__class__, "input_cost_metric"):
+            __class__.input_cost_metric = Counter(
+                'input_cost', 'Input cost'
+            )
+
+        if not hasattr(__class__, "output_cost_metric"):
+            __class__.output_cost_metric = Counter(
+                'output_cost', 'Output cost'
+            )
+
         input_queue = params.get("input_queue", default_input_queue)
         subscriber = params.get("subscriber", default_subscriber)
 
@@ -50,6 +70,9 @@ class Processor(Consumer):
         num_in = v.in_token
         num_out = v.out_token
 
+        __class__.input_token_metric.inc(num_in)
+        __class__.output_token_metric.inc(num_out)
+
         model_input_price, model_output_price = self.get_prices(price_list, modelname)
 
         if model_input_price == None:
@@ -59,6 +82,9 @@ class Processor(Consumer):
             cost_out = num_out * model_output_price
             cost_per_call = round(cost_in + cost_out, 6)
 
+            __class__.input_cost_metric.inc(cost_in)
+            __class__.output_cost_metric.inc(cost_out)
+
         print(f"Input Tokens: {num_in}", flush=True)
         print(f"Output Tokens: {num_out}", flush=True)
         print(f"Cost for call: ${cost_per_call}", flush=True)
@@ -72,4 +98,4 @@ class Processor(Consumer):
 
 def run():
 
-    Processor.start(module, __doc__)
\ No newline at end of file
+    Processor.start(module, __doc__)