From d02c7e4ea7fd76d3e801e4cc5c366ac7f44af241 Mon Sep 17 00:00:00 2001
From: Bartosz Sypytkowski <b.sypytkowski@gmail.com>
Date: Wed, 12 Jun 2024 08:27:37 +0200
Subject: [PATCH] feat: add workspace label to search open ai token metrics

---
 src/api/metrics.rs    | 19 +++++++++++++++----
 src/biz/search/ops.rs |  2 +-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/api/metrics.rs b/src/api/metrics.rs
index 53408ccd..c86f6e41 100644
--- a/src/api/metrics.rs
+++ b/src/api/metrics.rs
@@ -9,6 +9,7 @@ use prometheus_client::metrics::exemplar::CounterWithExemplar;
 use prometheus_client::metrics::family::Family;
 use prometheus_client::registry::Registry;
 use std::sync::Arc;
+use uuid::Uuid;
 
 pub fn metrics_scope() -> Scope {
   web::scope("/metrics").service(web::resource("").route(web::get().to(metrics_handler)))
@@ -38,6 +39,11 @@ pub struct ResultLabel {
   pub status_code: u16,
 }
 
+#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
+pub struct WorkspaceLabel {
+  pub workspace: String,
+}
+
 // Metrics contains list of metrics that are collected by the application.
 // Metric types: https://prometheus.io/docs/concepts/metric_types
 // Application handlers should call the corresponding methods to update the metrics.
@@ -46,7 +52,7 @@ pub struct RequestMetrics {
   requests_count: Family<PathLabel, Counter>,
   requests_latency: Family<PathLabel, CounterWithExemplar<TraceLabel>>,
   requests_result: Family<ResultLabel, CounterWithExemplar<TraceLabel>>,
-  openai_token_usage: Counter,
+  openai_token_usage: Family<WorkspaceLabel, Counter>,
 }
 
 #[derive(Clone, Hash, PartialEq, Eq, EncodeLabelSet, Debug, Default)]
@@ -60,7 +66,7 @@ impl RequestMetrics {
       requests_count: Family::default(),
       requests_latency: Family::default(),
       requests_result: Family::default(),
-      openai_token_usage: Counter::default(),
+      openai_token_usage: Family::default(),
     }
   }
 
@@ -91,8 +97,13 @@ impl RequestMetrics {
     af_metrics
   }
 
-  pub fn record_search_tokens_used(&self, tokens: u32) {
-    self.openai_token_usage.inc_by(tokens as u64);
+  pub fn record_search_tokens_used(&self, workspace_id: &Uuid, tokens: u32) {
+    self
+      .openai_token_usage
+      .get_or_create(&WorkspaceLabel {
+        workspace: workspace_id.to_string(),
+      })
+      .inc_by(tokens as u64);
   }
 
   // app services/middleware should call this method to increase the request count for the path
diff --git a/src/biz/search/ops.rs b/src/biz/search/ops.rs
index 3a1bd24a..96e024aa 100644
--- a/src/biz/search/ops.rs
+++ b/src/biz/search/ops.rs
@@ -33,7 +33,7 @@ pub async fn search_document(
     .map_err(|e| AppResponseError::new(ErrorCode::Internal, e.to_string()))?;
 
   let tokens_used = if let Some(usage) = embeddings.usage {
-    metrics.record_search_tokens_used(usage.total_tokens);
+    metrics.record_search_tokens_used(&workspace_id, usage.total_tokens);
     tracing::info!(
       "workspace {} OpenAI API search tokens used: {}",
       workspace_id,