From 0b93190cbd2be2f16b07cee4c59126cf8d5b7b89 Mon Sep 17 00:00:00 2001
From: Bartosz Sypytkowski <b.sypytkowski@gmail.com>
Date: Thu, 6 Jun 2024 06:38:58 +0200
Subject: [PATCH] chore: let prometheus track open ai tokens used on search

---
 src/api/metrics.rs    | 11 +++++++++++
 src/api/search.rs     |  3 ++-
 src/biz/search/ops.rs |  3 +++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/api/metrics.rs b/src/api/metrics.rs
index c70db095..53408ccd 100644
--- a/src/api/metrics.rs
+++ b/src/api/metrics.rs
@@ -46,6 +46,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,
 }
 
 #[derive(Clone, Hash, PartialEq, Eq, EncodeLabelSet, Debug, Default)]
@@ -59,6 +60,7 @@ impl RequestMetrics {
       requests_count: Family::default(),
       requests_latency: Family::default(),
       requests_result: Family::default(),
+      openai_token_usage: Counter::default(),
     }
   }
 
@@ -81,9 +83,18 @@ impl RequestMetrics {
       "status code of response",
       af_metrics.requests_result.clone(),
     );
+    af_registry.register(
+      "search_tokens_used",
+      "OpenAI API tokens used for search requests",
+      af_metrics.openai_token_usage.clone(),
+    );
     af_metrics
   }
 
+  pub fn record_search_tokens_used(&self, tokens: u32) {
+    self.openai_token_usage.inc_by(tokens as u64);
+  }
+
   // app services/middleware should call this method to increase the request count for the path
   pub fn record_request(&self, trace_id: Option<String>, path: String, ms: u64, status_code: u16) {
     self
diff --git a/src/api/search.rs b/src/api/search.rs
index 208e9cf2..56768469 100644
--- a/src/api/search.rs
+++ b/src/api/search.rs
@@ -37,6 +37,7 @@ async fn document_search(
       )
     },
   };
-  let resp = search_document(&state.pg_pool, openai, uid, workspace_id, request).await?;
+  let metrics = &*state.metrics.request_metrics;
+  let resp = search_document(&state.pg_pool, openai, uid, workspace_id, request, metrics).await?;
   Ok(AppResponse::Ok().with_data(resp).into())
 }
diff --git a/src/biz/search/ops.rs b/src/biz/search/ops.rs
index 38417457..3a1bd24a 100644
--- a/src/biz/search/ops.rs
+++ b/src/biz/search/ops.rs
@@ -1,3 +1,4 @@
+use crate::api::metrics::RequestMetrics;
 use app_error::ErrorCode;
 use database::index::{search_documents, SearchDocumentParams};
 use openai_dive::v1::models::EmbeddingsEngine;
@@ -17,6 +18,7 @@ pub async fn search_document(
   uid: i64,
   workspace_id: Uuid,
   request: SearchDocumentRequest,
+  metrics: &RequestMetrics,
 ) -> Result<Vec<SearchDocumentResponseItem>, AppResponseError> {
   let embeddings = openai
     .embeddings()
@@ -31,6 +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);
     tracing::info!(
       "workspace {} OpenAI API search tokens used: {}",
       workspace_id,