From 9ca904f943ff4f31daf4d25299c5ba33c37722c2 Mon Sep 17 00:00:00 2001
From: Bartosz Sypytkowski <b.sypytkowski@gmail.com>
Date: Tue, 18 Jun 2024 06:53:39 +0200
Subject: [PATCH] chore: change open ai token tracking to daily basis

---
 .../src/index/collab_embeddings_ops.rs        |  7 ++++--
 libs/database/src/index/search_ops.rs         |  8 ++++---
 .../20240618035048_af_workspace_ai_usage.sql  | 23 +++++++++++++++++++
 .../appflowy-indexer/src/collab_handle.rs     | 15 ++++++------
 4 files changed, 41 insertions(+), 12 deletions(-)
 create mode 100644 migrations/20240618035048_af_workspace_ai_usage.sql

diff --git a/libs/database/src/index/collab_embeddings_ops.rs b/libs/database/src/index/collab_embeddings_ops.rs
index 146ac90c..f1509974 100644
--- a/libs/database/src/index/collab_embeddings_ops.rs
+++ b/libs/database/src/index/collab_embeddings_ops.rs
@@ -41,8 +41,11 @@ pub async fn upsert_collab_embeddings(
   records: Vec<AFCollabEmbeddingParams>,
 ) -> Result<(), sqlx::Error> {
   if tokens_used > 0 {
-    sqlx::query(
-      "UPDATE af_workspace SET index_token_usage = index_token_usage + $2 WHERE workspace_id = $1",
+    sqlx::query(r#"
+      INSERT INTO af_workspace_ai_usage(created_at, workspace_id, search_requests, search_tokens_consumed, index_tokens_consumed)
+      VALUES (now()::date, $1, 0, 0, $2)
+      ON CONFLICT (created_at, workspace_id) DO UPDATE
+      SET index_tokens_consumed = af_workspace_ai_usage.index_tokens_consumed + $2"#,
     )
     .bind(workspace_id)
     .bind(tokens_used as i64)
diff --git a/libs/database/src/index/search_ops.rs b/libs/database/src/index/search_ops.rs
index 5afeb743..6c756230 100644
--- a/libs/database/src/index/search_ops.rs
+++ b/libs/database/src/index/search_ops.rs
@@ -13,9 +13,11 @@ pub async fn search_documents(
   let query = sqlx::query_as::<_, SearchDocumentItem>(
     r#"
     WITH workspace AS (
-      UPDATE af_workspace
-      SET search_token_usage = search_token_usage + $6
-      WHERE workspace_id = $2
+      INSERT INTO af_workspace_ai_usage(created_at, workspace_id, search_requests, search_tokens_consumed, index_tokens_consumed)
+      VALUES (now()::date, $2, 1, $6, 0)
+      ON CONFLICT (created_at, workspace_id) DO UPDATE
+      SET search_requests = af_workspace_ai_usage.search_requests + 1,
+          search_tokens_consumed = af_workspace_ai_usage.search_tokens_consumed + $6
       RETURNING workspace_id
     )
     SELECT
diff --git a/migrations/20240618035048_af_workspace_ai_usage.sql b/migrations/20240618035048_af_workspace_ai_usage.sql
new file mode 100644
index 00000000..9738a1dd
--- /dev/null
+++ b/migrations/20240618035048_af_workspace_ai_usage.sql
@@ -0,0 +1,23 @@
+CREATE TABLE IF NOT EXISTS af_workspace_ai_usage (
+    created_at DATE NOT NULL,               -- day level of granularity
+    workspace_id UUID NOT NULL,             -- workspace id for which the usage is being recorded
+    search_requests INT,                    -- number of search requests made
+    search_tokens_consumed BIGINT,          -- number of tokens consumed for search requests
+    index_tokens_consumed BIGINT,           -- number of tokens consumed for indexing documents
+    PRIMARY KEY (created_at, workspace_id)
+);
+
+-- migrate token usage data from af_workspace to af_workspace_ai_usage
+INSERT INTO af_workspace_ai_usage (created_at, workspace_id, search_tokens_consumed, index_tokens_consumed)
+SELECT
+    now()::date as created_at,
+    workspace_id,
+    search_token_usage as search_tokens_consumed,
+    index_token_usage as index_tokens_consumed
+FROM af_workspace
+WHERE search_token_usage IS NOT NULL
+   OR index_token_usage IS NOT NULL;
+
+-- drop the redundant columns from af_workspace
+ALTER TABLE af_workspace DROP COLUMN IF EXISTS search_token_usage;
+ALTER TABLE af_workspace DROP COLUMN IF EXISTS index_token_usage;
\ No newline at end of file
diff --git a/services/appflowy-indexer/src/collab_handle.rs b/services/appflowy-indexer/src/collab_handle.rs
index c33b1017..4b67489d 100644
--- a/services/appflowy-indexer/src/collab_handle.rs
+++ b/services/appflowy-indexer/src/collab_handle.rs
@@ -363,13 +363,14 @@ mod test {
 
     assert_eq!(contents.len(), 1);
 
-    let tokens: i64 =
-      sqlx::query("SELECT index_token_usage from af_workspace WHERE workspace_id = $1")
-        .bind(workspace_id)
-        .fetch_one(&db)
-        .await
-        .unwrap()
-        .get(0);
+    let tokens: i64 = sqlx::query(
+      "SELECT index_tokens_consumed from af_workspace_ai_usage WHERE workspace_id = $1",
+    )
+    .bind(workspace_id)
+    .fetch_one(&db)
+    .await
+    .unwrap()
+    .get(0);
     assert_ne!(tokens, 0);
   }
 }