From 9ddc84e6cf6065e6a9946090717eac6867a6f5cb Mon Sep 17 00:00:00 2001 From: KernelDeimos Date: Wed, 27 Nov 2024 13:53:22 -0500 Subject: [PATCH] dev: record token-based usages --- .../src/modules/puterai/AIChatService.js | 46 ++++++++++++++++--- .../database/BaseDatabaseAccessService.js | 15 ++++++ .../database/SqliteDatabaseAccessService.js | 3 ++ .../database/sqlite_setup/0033_ai-usage.sql | 34 ++++++++++++++ 4 files changed, 91 insertions(+), 7 deletions(-) create mode 100644 src/backend/src/services/database/sqlite_setup/0033_ai-usage.sql diff --git a/src/backend/src/modules/puterai/AIChatService.js b/src/backend/src/modules/puterai/AIChatService.js index 24df4b95..f47b299a 100644 --- a/src/backend/src/modules/puterai/AIChatService.js +++ b/src/backend/src/modules/puterai/AIChatService.js @@ -1,4 +1,5 @@ const BaseService = require("../../services/BaseService"); +const { DB_WRITE } = require("../../services/database/consts"); const { TypeSpec } = require("../../services/drivers/meta/Construct"); const { TypedValue } = require("../../services/drivers/meta/Runtime"); const { Context } = require("../../util/context"); @@ -21,16 +22,47 @@ class AIChatService extends BaseService { _init () { this.kvkey = this.modules.uuidv4(); + this.db = this.services.get('database').get(DB_WRITE, 'ai-usage'); + const svc_event = this.services.get('event'); svc_event.on('ai.prompt.report-usage', async (_, details) => { - const user_id = details.actor?.type?.user?.id; - const app_id = details.actor?.type?.app?.id; - const service_name = details.service_used; - const model_name = details.model_used; - const value_uint_1 = details.usage?.input_tokens; - const value_uint_2 = details.usage?.output_tokens; + const values = { + user_id: details.actor?.type?.user?.id, + app_id: details.actor?.type?.app?.id, + service_name: details.service_used, + model_name: details.model_used, + value_uint_1: details.usage?.input_tokens, + value_uint_2: details.usage?.output_tokens, + }; - console.log('reporting usage', { user_id, app_id, service_name, model_name, value_uint_1, value_uint_2 }); }); + let model_details = this.detail_model_map[details.model_used]; + if ( Array.isArray(model_details) ) { + for ( const model of model_details ) { + if ( model.provider === details.service_used ) { + model_details = model; + break; + } + } + } + if ( Array.isArray(model_details) ) { + model_details = model_details[0]; + } + if ( model_details ) { + values.cost = 0 + // for formatting + + model_details.cost.input * details.usage.input_tokens + // cents/MTok tokens + + + + model_details.cost.output * details.usage.output_tokens + // cents/MTok tokens + ; + } else { + this.log.error('could not find model details', { details }); + } + + await this.db.insert('ai_usage', values); + }); } async ['__on_boot.consolidation'] () { diff --git a/src/backend/src/services/database/BaseDatabaseAccessService.js b/src/backend/src/services/database/BaseDatabaseAccessService.js index ebac058b..239262c5 100644 --- a/src/backend/src/services/database/BaseDatabaseAccessService.js +++ b/src/backend/src/services/database/BaseDatabaseAccessService.js @@ -52,6 +52,21 @@ class BaseDatabaseAccessService extends BaseService { return this._write(query, params); } + insert (table_name, data) { + const values = Object.values(data); + const sql = this._gen_insert_sql(table_name, data); + console.log('INSERT SQL', sql); + return this.write(sql, values); + } + + _gen_insert_sql (table_name, data) { + const cols = Object.keys(data); + return 'INSERT INTO `' + table_name + '` ' + + '(' + cols.map(str => '`' + str + '`').join(', ') + ') ' + + 'VALUES (' + cols.map(() => '?').join(', ') + ')'; + } + + batch_write (statements) { return this._batch_write(statements); } diff --git a/src/backend/src/services/database/SqliteDatabaseAccessService.js b/src/backend/src/services/database/SqliteDatabaseAccessService.js index 8d67fbe4..d32ae224 100644 --- a/src/backend/src/services/database/SqliteDatabaseAccessService.js +++ b/src/backend/src/services/database/SqliteDatabaseAccessService.js @@ -141,6 +141,9 @@ class SqliteDatabaseAccessService extends BaseDatabaseAccessService { [29, [ '0032_signup_metadata.sql', ]], + [30, [ + '0033_ai-usage.sql', + ]], ]; // Database upgrade logic diff --git a/src/backend/src/services/database/sqlite_setup/0033_ai-usage.sql b/src/backend/src/services/database/sqlite_setup/0033_ai-usage.sql new file mode 100644 index 00000000..c3521c85 --- /dev/null +++ b/src/backend/src/services/database/sqlite_setup/0033_ai-usage.sql @@ -0,0 +1,34 @@ +CREATE TABLE `ai_usage` ( + `id` INTEGER PRIMARY KEY AUTOINCREMENT, + `user_id` INTEGER NOT NULL, + `app_id` INTEGER DEFAULT NULL, + `service_name` TEXT NOT NULL, + `model_name` TEXT NOT NULL, + + -- set this to a string when service:model alone does not make + -- the numeric values below fungible + `price_modifier` TEXT DEFAULT NULL, + + -- expected cost of request in ยตยข (microcents) + `cost` int DEFAULT NULL, + + -- input tokens + `value_uint_1` int DEFAULT NULL, + -- output tokens + `value_uint_2` int DEFAULT NULL, + + -- miscelaneous values for future use + `value_uint_3` int DEFAULT NULL, + `value_uint_4` int DEFAULT NULL, + `value_uint_5` int DEFAULT NULL, + + `created_at` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + + FOREIGN KEY("user_id") REFERENCES "user" ("id") ON DELETE CASCADE ON UPDATE CASCADE, + FOREIGN KEY("app_id") REFERENCES "apps" ("id") ON DELETE SET NULL ON UPDATE CASCADE +); + +CREATE INDEX `idx_ai_usage_service_name` ON `ai_usage` (`service_name`); +CREATE INDEX `idx_ai_usage_model_name` ON `ai_usage` (`model_name`); +CREATE INDEX `idx_ai_usage_price_modifier` ON `ai_usage` (`price_modifier`); +CREATE INDEX `idx_ai_usage_created_at` ON `ai_usage` (`created_at`);