github-mirrors/AppFlowy-Cloud
1
0
Fork 0
mirror of https://github.com/AppFlowy-IO/AppFlowy-Cloud.git synced 2025-04-19 03:24:42 -04:00
Code Issues Projects Releases Packages Wiki Activity

chore: change open ai token tracking to daily basis

Browse source
This commit is contained in:
Bartosz Sypytkowski 2024-06-18 06:53:39 +02:00
parent ab01363786
commit 9ca904f943
4 changed files with 41 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
libs/database/src/index/collab_embeddings_ops.rs
View file

@ -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)

8
libs/database/src/index/search_ops.rs
View file

@ -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

23
migrations/20240618035048_af_workspace_ai_usage.sql Normal file
View file

@ -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;

15
services/appflowy-indexer/src/collab_handle.rs
View file

@ -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);
}
}