From e990bad68f2b28e28d9a6853eaf6ceabdcbcf273 Mon Sep 17 00:00:00 2001
From: "Nathan.fooo" <86001920+appflowy@users.noreply.github.com>
Date: Mon, 20 Jan 2025 23:22:33 +0800
Subject: [PATCH] chore: add metrics for ai (#1184)

---
 src/api/ai.rs           |  3 +++
 src/biz/chat/metrics.rs | 30 ++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/src/api/ai.rs b/src/api/ai.rs
index 8d7c7a347..484f75e9b 100644
--- a/src/api/ai.rs
+++ b/src/api/ai.rs
@@ -37,6 +37,7 @@ async fn stream_complete_text_handler(
 ) -> actix_web::Result<HttpResponse> {
   let ai_model = ai_model_from_header(&req);
   let params = payload.into_inner();
+  state.metrics.ai_metrics.record_total_completion_count(1);
 
   match state
     .ai_client
@@ -80,6 +81,7 @@ async fn summarize_row_handler(
         );
       }
 
+      state.metrics.ai_metrics.record_total_summary_row_count(1);
       let ai_model = ai_model_from_header(&req);
       let result = state.ai_client.summarize_row(&content, ai_model).await;
       let resp = match result {
@@ -105,6 +107,7 @@ async fn translate_row_handler(
 ) -> actix_web::Result<Json<AppResponse<TranslateRowResponse>>> {
   let params = payload.into_inner();
   let ai_model = ai_model_from_header(&req);
+  state.metrics.ai_metrics.record_total_translate_row_count(1);
   match state.ai_client.translate_row(params.data, ai_model).await {
     Ok(resp) => Ok(AppResponse::Ok().with_data(resp).into()),
     Err(err) => {
diff --git a/src/biz/chat/metrics.rs b/src/biz/chat/metrics.rs
index 3d6fdcf17..69c0f8534 100644
--- a/src/biz/chat/metrics.rs
+++ b/src/biz/chat/metrics.rs
@@ -5,6 +5,9 @@ pub struct AIMetrics {
   total_stream_count: Counter,
   failed_stream_count: Counter,
   stream_image_count: Counter,
+  total_completion_count: Counter,
+  total_summary_row_count: Counter,
+  total_translate_row_count: Counter,
 }
 
 impl AIMetrics {
@@ -28,6 +31,21 @@ impl AIMetrics {
       "Total count of image streams processed",
       metrics.stream_image_count.clone(),
     );
+    realtime_registry.register(
+      "total_completion_count",
+      "Total count of completions processed",
+      metrics.total_completion_count.clone(),
+    );
+    realtime_registry.register(
+      "total_summary_row_count",
+      "Total count of summary rows processed",
+      metrics.total_summary_row_count.clone(),
+    );
+    realtime_registry.register(
+      "total_translate_row_count",
+      "Total count of translation rows processed",
+      metrics.total_translate_row_count.clone(),
+    );
 
     metrics
   }
@@ -43,4 +61,16 @@ impl AIMetrics {
   pub fn record_stream_image_count(&self, count: u64) {
     self.stream_image_count.inc_by(count);
   }
+
+  pub fn record_total_completion_count(&self, count: u64) {
+    self.total_completion_count.inc_by(count);
+  }
+
+  pub fn record_total_summary_row_count(&self, count: u64) {
+    self.total_summary_row_count.inc_by(count);
+  }
+
+  pub fn record_total_translate_row_count(&self, count: u64) {
+    self.total_translate_row_count.inc_by(count);
+  }
 }