diff --git a/src/ipc/database/cloudHandler.ts b/src/ipc/database/cloudHandler.ts index cb97d69..fce86d9 100644 --- a/src/ipc/database/cloudHandler.ts +++ b/src/ipc/database/cloudHandler.ts @@ -126,6 +126,34 @@ function ensureDatabaseInitialized(dbPath: string): void { value TEXT NOT NULL ); `); + + // Create token_usage table + db.exec(` + CREATE TABLE IF NOT EXISTS token_usage ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + account_id TEXT NOT NULL, + model TEXT NOT NULL, + prompt_tokens INTEGER NOT NULL DEFAULT 0, + completion_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + timestamp INTEGER NOT NULL, + request_type TEXT + ); + `); + + // Migration: Check if request_type column exists on token_usage + const tokenUsageTableInfoRaw = db.pragma('table_info(token_usage)') as any[]; + const tokenUsageTableInfo = parseRows(TableInfoRowSchema, tokenUsageTableInfoRaw, 'cloud.token_usage.tableInfo'); + const hasRequestType = tokenUsageTableInfo.some((col) => col.name === 'request_type'); + if (!hasRequestType) { + db.exec('ALTER TABLE token_usage ADD COLUMN request_type TEXT'); + } + + // Create indexes for token_usage + db.exec(`CREATE INDEX IF NOT EXISTS idx_token_usage_account ON token_usage(account_id);`); + db.exec(`CREATE INDEX IF NOT EXISTS idx_token_usage_timestamp ON token_usage(timestamp);`); + db.exec(`CREATE INDEX IF NOT EXISTS idx_token_usage_model ON token_usage(model);`); + db.exec(`CREATE INDEX IF NOT EXISTS idx_token_usage_account_timestamp ON token_usage(account_id, timestamp);`); } catch (error) { logger.error('Failed to initialize cloud database schema', error); throw error; @@ -137,7 +165,7 @@ function ensureDatabaseInitialized(dbPath: string): void { /** * Gets a connection to the cloud accounts database. */ -function getCloudDb(): { +export function getCloudDb(): { raw: Database.Database; orm: BetterSQLite3Database; } { diff --git a/src/ipc/database/schema.ts b/src/ipc/database/schema.ts index f411204..b80b38b 100644 --- a/src/ipc/database/schema.ts +++ b/src/ipc/database/schema.ts @@ -27,3 +27,14 @@ export const itemTable = sqliteTable('ItemTable', { key: text('key').primaryKey(), value: text('value'), }); + +export const tokenUsage = sqliteTable('token_usage', { + id: integer('id').primaryKey({ autoIncrement: true }), + accountId: text('account_id').notNull(), + model: text('model').notNull(), + promptTokens: integer('prompt_tokens').notNull().default(0), + completionTokens: integer('completion_tokens').notNull().default(0), + totalTokens: integer('total_tokens').notNull().default(0), + timestamp: integer('timestamp').notNull(), + requestType: text('request_type'), +}); diff --git a/src/ipc/database/usageHandler.ts b/src/ipc/database/usageHandler.ts new file mode 100644 index 0000000..7403603 --- /dev/null +++ b/src/ipc/database/usageHandler.ts @@ -0,0 +1,244 @@ +import { logger } from '../../utils/logger'; +import { getCloudDb } from './cloudHandler'; +import { tokenUsage } from './schema'; +import { + HourlyUsageSchema, + DailyUsageSchema, + WeeklyUsageSchema, + MonthlyUsageSchema, + ModelUsageSchema, +} from '../../types/db'; + +export interface RecordUsageParams { + accountId: string; + model: string; + promptTokens: number; + completionTokens: number; + totalTokens: number; + timestamp: number; + requestType?: string; +} + +function buildDateFilter(start?: number, end?: number): string { + const conditions: string[] = []; + if (start !== undefined) { + conditions.push(`timestamp >= ${start}`); + } + if (end !== undefined) { + conditions.push(`timestamp <= ${end}`); + } + return conditions.length > 0 ? `AND ${conditions.join(' AND ')}` : ''; +} + +function buildAccountFilter(accountId?: string): string { + return accountId ? `AND account_id = '${accountId}'` : ''; +} + +export class TokenUsageRepo { + static recordUsage(params: RecordUsageParams): void { + const { raw, orm } = getCloudDb(); + try { + orm + .insert(tokenUsage) + .values({ + accountId: params.accountId, + model: params.model, + promptTokens: params.promptTokens, + completionTokens: params.completionTokens, + totalTokens: params.totalTokens, + timestamp: params.timestamp, + requestType: params.requestType ?? null, + }) + .run(); + } catch (error) { + logger.error('Failed to record token usage', error); + } finally { + raw.close(); + } + } + + static getUsageByHour( + accountId?: string, + start?: number, + end?: number, + ): Array<{ + bucket: string; + promptTokens: number; + completionTokens: number; + totalTokens: number; + requests: number; + }> { + const { raw } = getCloudDb(); + try { + const sql = ` + SELECT + strftime('%Y-%m-%d %H:00', timestamp / 1000, 'unixepoch', 'localtime') as bucket, + SUM(prompt_tokens) as promptTokens, + SUM(completion_tokens) as completionTokens, + SUM(total_tokens) as totalTokens, + COUNT(*) as requests + FROM token_usage + WHERE 1=1 + ${buildAccountFilter(accountId)} + ${buildDateFilter(start, end)} + GROUP BY bucket + ORDER BY bucket + `; + const rows = raw.prepare(sql).all() as unknown[]; + return rows.map((row) => HourlyUsageSchema.parse(row)); + } catch (error) { + logger.error('Failed to get hourly usage', error); + return []; + } finally { + raw.close(); + } + } + + static getUsageByDay( + accountId?: string, + start?: number, + end?: number, + ): Array<{ + bucket: string; + promptTokens: number; + completionTokens: number; + totalTokens: number; + requests: number; + }> { + const { raw } = getCloudDb(); + try { + const sql = ` + SELECT + strftime('%Y-%m-%d', timestamp / 1000, 'unixepoch', 'localtime') as bucket, + SUM(prompt_tokens) as promptTokens, + SUM(completion_tokens) as completionTokens, + SUM(total_tokens) as totalTokens, + COUNT(*) as requests + FROM token_usage + WHERE 1=1 + ${buildAccountFilter(accountId)} + ${buildDateFilter(start, end)} + GROUP BY bucket + ORDER BY bucket + `; + const rows = raw.prepare(sql).all() as unknown[]; + return rows.map((row) => DailyUsageSchema.parse(row)); + } catch (error) { + logger.error('Failed to get daily usage', error); + return []; + } finally { + raw.close(); + } + } + + static getUsageByWeek( + accountId?: string, + start?: number, + end?: number, + ): Array<{ + bucket: string; + promptTokens: number; + completionTokens: number; + totalTokens: number; + requests: number; + }> { + const { raw } = getCloudDb(); + try { + const sql = ` + SELECT + strftime('%Y-%W', timestamp / 1000, 'unixepoch', 'localtime') as bucket, + SUM(prompt_tokens) as promptTokens, + SUM(completion_tokens) as completionTokens, + SUM(total_tokens) as totalTokens, + COUNT(*) as requests + FROM token_usage + WHERE 1=1 + ${buildAccountFilter(accountId)} + ${buildDateFilter(start, end)} + GROUP BY bucket + ORDER BY bucket + `; + const rows = raw.prepare(sql).all() as unknown[]; + return rows.map((row) => WeeklyUsageSchema.parse(row)); + } catch (error) { + logger.error('Failed to get weekly usage', error); + return []; + } finally { + raw.close(); + } + } + + static getUsageByMonth( + accountId?: string, + start?: number, + end?: number, + ): Array<{ + bucket: string; + promptTokens: number; + completionTokens: number; + totalTokens: number; + requests: number; + }> { + const { raw } = getCloudDb(); + try { + const sql = ` + SELECT + strftime('%Y-%m', timestamp / 1000, 'unixepoch', 'localtime') as bucket, + SUM(prompt_tokens) as promptTokens, + SUM(completion_tokens) as completionTokens, + SUM(total_tokens) as totalTokens, + COUNT(*) as requests + FROM token_usage + WHERE 1=1 + ${buildAccountFilter(accountId)} + ${buildDateFilter(start, end)} + GROUP BY bucket + ORDER BY bucket + `; + const rows = raw.prepare(sql).all() as unknown[]; + return rows.map((row) => MonthlyUsageSchema.parse(row)); + } catch (error) { + logger.error('Failed to get monthly usage', error); + return []; + } finally { + raw.close(); + } + } + + static getUsageByModel( + accountId?: string, + start?: number, + end?: number, + ): Array<{ + model: string; + promptTokens: number; + completionTokens: number; + totalTokens: number; + requests: number; + }> { + const { raw } = getCloudDb(); + try { + const sql = ` + SELECT + model, + SUM(prompt_tokens) as promptTokens, + SUM(completion_tokens) as completionTokens, + SUM(total_tokens) as totalTokens, + COUNT(*) as requests + FROM token_usage + WHERE 1=1 + ${buildAccountFilter(accountId)} + ${buildDateFilter(start, end)} + GROUP BY model + ORDER BY totalTokens DESC + `; + const rows = raw.prepare(sql).all() as unknown[]; + return rows.map((row) => ModelUsageSchema.parse(row)); + } catch (error) { + logger.error('Failed to get model usage', error); + return []; + } finally { + raw.close(); + } + } +} diff --git a/src/types/db.ts b/src/types/db.ts index 30b9efa..6073d77 100644 --- a/src/types/db.ts +++ b/src/types/db.ts @@ -70,3 +70,66 @@ export const TableInfoRowSchema = z.object({ }); export type TableInfoRow = z.infer; + +export const TokenUsageRowSchema = z.object({ + id: z.number(), + account_id: z.string(), + model: z.string(), + prompt_tokens: z.number(), + completion_tokens: z.number(), + total_tokens: z.number(), + timestamp: z.number(), + request_type: z.string().nullable(), +}); + +export type TokenUsageRow = z.infer; + +export const HourlyUsageSchema = z.object({ + bucket: z.string(), + promptTokens: z.number(), + completionTokens: z.number(), + totalTokens: z.number(), + requests: z.number(), +}); + +export type HourlyUsage = z.infer; + +export const DailyUsageSchema = z.object({ + bucket: z.string(), + promptTokens: z.number(), + completionTokens: z.number(), + totalTokens: z.number(), + requests: z.number(), +}); + +export type DailyUsage = z.infer; + +export const WeeklyUsageSchema = z.object({ + bucket: z.string(), + promptTokens: z.number(), + completionTokens: z.number(), + totalTokens: z.number(), + requests: z.number(), +}); + +export type WeeklyUsage = z.infer; + +export const MonthlyUsageSchema = z.object({ + bucket: z.string(), + promptTokens: z.number(), + completionTokens: z.number(), + totalTokens: z.number(), + requests: z.number(), +}); + +export type MonthlyUsage = z.infer; + +export const ModelUsageSchema = z.object({ + model: z.string(), + promptTokens: z.number(), + completionTokens: z.number(), + totalTokens: z.number(), + requests: z.number(), +}); + +export type ModelUsage = z.infer;