chore: let prometheus track open ai tokens used on search

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

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())
 }

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,