From 45b327dcc8c85824f3ae130148bb993fd9480217 Mon Sep 17 00:00:00 2001
From: Nikhil Sinha <nikhil.sinha@parseable.com>
Date: Sun, 16 Nov 2025 07:33:14 -0800
Subject: [PATCH] add cached and reasoning token to the metrics

to be used in enterprise
---
 src/metrics/mod.rs | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/src/metrics/mod.rs b/src/metrics/mod.rs
index 1a5f177dc..3d1b492bd 100644
--- a/src/metrics/mod.rs
+++ b/src/metrics/mod.rs
@@ -334,6 +334,30 @@ pub static TOTAL_OUTPUT_LLM_TOKENS_BY_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
     .expect("metric can be created")
 });
 
+pub static TOTAL_CACHED_LLM_TOKENS_BY_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
+        Opts::new(
+            "total_cached_llm_tokens_by_date",
+            "Total cached LLM tokens used by date",
+        )
+        .namespace(METRICS_NAMESPACE),
+        &["provider", "model", "date"],
+    )
+    .expect("metric can be created")
+});
+
+pub static TOTAL_REASONING_LLM_TOKENS_BY_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
+        Opts::new(
+            "total_reasoning_llm_tokens_by_date",
+            "Total reasoning LLM tokens used by date",
+        )
+        .namespace(METRICS_NAMESPACE),
+        &["provider", "model", "date"],
+    )
+    .expect("metric can be created")
+});
+
 pub static STORAGE_REQUEST_RESPONSE_TIME: Lazy<HistogramVec> = Lazy::new(|| {
     HistogramVec::new(
         HistogramOpts::new("storage_request_response_time", "Storage Request Latency")
@@ -433,6 +457,12 @@ fn custom_metrics(registry: &Registry) {
     registry
         .register(Box::new(TOTAL_OUTPUT_LLM_TOKENS_BY_DATE.clone()))
         .expect("metric can be registered");
+    registry
+        .register(Box::new(TOTAL_CACHED_LLM_TOKENS_BY_DATE.clone()))
+        .expect("metric can be registered");
+    registry
+        .register(Box::new(TOTAL_REASONING_LLM_TOKENS_BY_DATE.clone()))
+        .expect("metric can be registered");
     registry
         .register(Box::new(STORAGE_REQUEST_RESPONSE_TIME.clone()))
         .expect("metric can be registered");
@@ -566,6 +596,23 @@ pub fn increment_output_llm_tokens_by_date(provider: &str, model: &str, tokens:
         .inc_by(tokens);
 }
 
+pub fn increment_cached_llm_tokens_by_date(provider: &str, model: &str, tokens: u64, date: &str) {
+    TOTAL_CACHED_LLM_TOKENS_BY_DATE
+        .with_label_values(&[provider, model, date])
+        .inc_by(tokens);
+}
+
+pub fn increment_reasoning_llm_tokens_by_date(
+    provider: &str,
+    model: &str,
+    tokens: u64,
+    date: &str,
+) {
+    TOTAL_REASONING_LLM_TOKENS_BY_DATE
+        .with_label_values(&[provider, model, date])
+        .inc_by(tokens);
+}
+
 use actix_web::HttpResponse;
 
 pub async fn get() -> Result<impl Responder, MetricsError> {