chore: change open ai token tracking to daily basis

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(
+    sqlx::query(r#"
-      "UPDATE af_workspace SET index_token_usage = index_token_usage + $2 WHERE workspace_id = $1",
+      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)

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
+      INSERT INTO af_workspace_ai_usage(created_at, workspace_id, search_requests, search_tokens_consumed, index_tokens_consumed)
-      SET search_token_usage = search_token_usage + $6
+      VALUES (now()::date, $2, 1, $6, 0)
-      WHERE workspace_id = $2
+      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

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;

services/appflowy-indexer/src/collab_handle.rs

@@ -363,8 +363,9 @@ mod test {
 
     assert_eq!(contents.len(), 1);
 
-    let tokens: i64 =
+    let tokens: i64 = sqlx::query(
-      sqlx::query("SELECT index_token_usage from af_workspace WHERE workspace_id = $1")
+      "SELECT index_tokens_consumed from af_workspace_ai_usage WHERE workspace_id = $1",
+    )
     .bind(workspace_id)
     .fetch_one(&db)
     .await