feat: add workspace label to search open ai token metrics

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

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,