From 9ca904f943ff4f31daf4d25299c5ba33c37722c2 Mon Sep 17 00:00:00 2001 From: Bartosz Sypytkowski Date: Tue, 18 Jun 2024 06:53:39 +0200 Subject: [PATCH 1/3] 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, ) -> 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); } } From a5bca1852ba397da7fda9d15972006eb4247ef8d Mon Sep 17 00:00:00 2001 From: Bartosz Sypytkowski Date: Tue, 18 Jun 2024 08:53:02 +0200 Subject: [PATCH 2/3] chore: change test query to macro --- services/appflowy-indexer/src/collab_handle.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/services/appflowy-indexer/src/collab_handle.rs b/services/appflowy-indexer/src/collab_handle.rs index 4b67489d..28ac33d6 100644 --- a/services/appflowy-indexer/src/collab_handle.rs +++ b/services/appflowy-indexer/src/collab_handle.rs @@ -363,14 +363,14 @@ mod test { assert_eq!(contents.len(), 1); - let tokens: i64 = sqlx::query( + let tokens = sqlx::query_scalar!( "SELECT index_tokens_consumed from af_workspace_ai_usage WHERE workspace_id = $1", + workspace_id ) - .bind(workspace_id) .fetch_one(&db) .await .unwrap() - .get(0); + .unwrap_or(0); assert_ne!(tokens, 0); } } From f11703949c0d73b22989ab4558898ab4205d500c Mon Sep 17 00:00:00 2001 From: Bartosz Sypytkowski Date: Tue, 18 Jun 2024 09:34:10 +0200 Subject: [PATCH 3/3] chore: turn off sqlx macros in tests --- services/appflowy-indexer/src/collab_handle.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/services/appflowy-indexer/src/collab_handle.rs b/services/appflowy-indexer/src/collab_handle.rs index 28ac33d6..daa137e5 100644 --- a/services/appflowy-indexer/src/collab_handle.rs +++ b/services/appflowy-indexer/src/collab_handle.rs @@ -287,7 +287,7 @@ mod test { use collab::preclude::Collab; use collab_document::document::Document; use collab_entity::CollabType; - use sqlx::Row; + use sqlx::{Postgres, Row}; use workspace_template::document::get_started::get_started_document_data; @@ -363,10 +363,10 @@ mod test { assert_eq!(contents.len(), 1); - let tokens = sqlx::query_scalar!( + let tokens: i64 = sqlx::query_scalar::>( "SELECT index_tokens_consumed from af_workspace_ai_usage WHERE workspace_id = $1", - workspace_id ) + .bind(workspace_id) .fetch_one(&db) .await .unwrap()