@theihtisham/budget-llm 1.0.0

package/src/database.ts

@@ -0,0 +1,496 @@
+import initSqlJs, { Database as SqlJsDatabase } from 'sql.js';
+import fs from 'fs';
+import path from 'path';
+import { getDbPath, log } from './config';
+import type { CostRecord, CostSummary, ProviderSummary, ModelSummary, BudgetStatus, BudgetConfig, CacheEntry, DashboardData, ProviderId } from './types';
+
+let db: SqlJsDatabase | null = null;
+let dbPath: string = '';
+let saveTimer: ReturnType<typeof setTimeout> | null = null;
+
+/**
+ * Get or initialize the database. sql.js is async so we init once.
+ */
+export async function initDb(): Promise<SqlJsDatabase> {
+  if (db) return db;
+
+  dbPath = getDbPath();
+  const dir = path.dirname(dbPath);
+
+  if (!fs.existsSync(dir)) {
+    fs.mkdirSync(dir, { recursive: true });
+  }
+
+  const SQL = await initSqlJs();
+
+  // Load existing database if it exists
+  if (fs.existsSync(dbPath)) {
+    const fileBuffer = fs.readFileSync(dbPath);
+    db = new SQL.Database(fileBuffer);
+  } else {
+    db = new SQL.Database();
+  }
+
+  migrate(db);
+  log.info('Database initialized and migrated');
+
+  // Auto-save every 5 seconds if dirty
+  saveTimer = setInterval(() => {
+    saveDb();
+  }, 5000).unref();
+
+  return db;
+}
+
+/**
+ * Synchronous access — throws if not initialized yet.
+ * Call initDb() during server startup before using this.
+ */
+export function getDb(): SqlJsDatabase {
+  if (!db) {
+    throw new Error('Database not initialized. Call initDb() first.');
+  }
+  return db;
+}
+
+function migrate(database: SqlJsDatabase): void {
+  database.run(`
+    CREATE TABLE IF NOT EXISTS cost_records (
+      id TEXT PRIMARY KEY,
+      request_id TEXT NOT NULL,
+      timestamp TEXT NOT NULL,
+      provider TEXT NOT NULL,
+      model TEXT NOT NULL,
+      task_type TEXT NOT NULL,
+      input_tokens INTEGER NOT NULL,
+      output_tokens INTEGER NOT NULL,
+      input_cost REAL NOT NULL,
+      output_cost REAL NOT NULL,
+      total_cost REAL NOT NULL,
+      latency_ms INTEGER NOT NULL,
+      cached INTEGER NOT NULL DEFAULT 0
+    );
+  `);
+  database.run(`CREATE INDEX IF NOT EXISTS idx_cost_timestamp ON cost_records(timestamp);`);
+  database.run(`CREATE INDEX IF NOT EXISTS idx_cost_provider ON cost_records(provider);`);
+  database.run(`CREATE INDEX IF NOT EXISTS idx_cost_model ON cost_records(model);`);
+
+  database.run(`
+    CREATE TABLE IF NOT EXISTS cache_entries (
+      prompt_hash TEXT PRIMARY KEY,
+      model TEXT NOT NULL,
+      provider TEXT NOT NULL,
+      response TEXT NOT NULL,
+      input_tokens INTEGER NOT NULL,
+      output_tokens INTEGER NOT NULL,
+      cost REAL NOT NULL,
+      created_at INTEGER NOT NULL,
+      expires_at INTEGER NOT NULL,
+      hit_count INTEGER NOT NULL DEFAULT 0
+    );
+  `);
+  database.run(`CREATE INDEX IF NOT EXISTS idx_cache_expires ON cache_entries(expires_at);`);
+
+  database.run(`
+    CREATE TABLE IF NOT EXISTS budget_config (
+      key TEXT PRIMARY KEY,
+      value TEXT NOT NULL
+    );
+  `);
+
+  database.run(`
+    CREATE TABLE IF NOT EXISTS api_keys (
+      provider TEXT PRIMARY KEY,
+      encrypted_key TEXT NOT NULL,
+      updated_at TEXT NOT NULL
+    );
+  `);
+
+  log.info('Database initialized and migrated');
+}
+
+function saveDb(): void {
+  if (!db || !dbPath) return;
+  try {
+    const data = db.export();
+    const buffer = Buffer.from(data);
+    fs.writeFileSync(dbPath, buffer);
+  } catch (err) {
+    log.error('Failed to save database', err);
+  }
+}
+
+// ---- Helper to query single row ----
+
+function queryOne<T>(sql: string, params: (string | number | null)[] = []): T | null {
+  const database = getDb();
+  const stmt = database.prepare(sql);
+  stmt.bind(params);
+  try {
+    if (stmt.step()) {
+      const row = stmt.getAsObject() as T;
+      return row;
+    }
+    return null;
+  } finally {
+    stmt.free();
+  }
+}
+
+function queryAll<T>(sql: string, params: (string | number | null)[] = []): T[] {
+  const database = getDb();
+  const stmt = database.prepare(sql);
+  stmt.bind(params);
+  const results: T[] = [];
+  try {
+    while (stmt.step()) {
+      results.push(stmt.getAsObject() as T);
+    }
+    return results;
+  } finally {
+    stmt.free();
+  }
+}
+
+// ---- Cost Record Operations ----
+
+export function insertCostRecord(record: CostRecord): void {
+  const database = getDb();
+  database.run(
+    `INSERT INTO cost_records (id, request_id, timestamp, provider, model, task_type,
+      input_tokens, output_tokens, input_cost, output_cost, total_cost, latency_ms, cached)
+    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+    [
+      record.id,
+      record.requestId,
+      record.timestamp,
+      record.provider,
+      record.model,
+      record.taskType,
+      record.inputTokens,
+      record.outputTokens,
+      record.inputCost,
+      record.outputCost,
+      record.totalCost,
+      record.latencyMs,
+      record.cached ? 1 : 0,
+    ]
+  );
+  saveDb();
+}
+
+export function getDailySpend(): number {
+  const todayStr = new Date().toISOString().slice(0, 10);
+  const row = queryOne<{ spent: number }>(
+    `SELECT COALESCE(SUM(total_cost), 0) as spent FROM cost_records WHERE timestamp >= ?`,
+    [todayStr]
+  );
+  return row?.spent ?? 0;
+}
+
+export function getMonthlySpend(): number {
+  const now = new Date();
+  const monthStart = new Date(now.getFullYear(), now.getMonth(), 1).toISOString();
+  const row = queryOne<{ spent: number }>(
+    `SELECT COALESCE(SUM(total_cost), 0) as spent FROM cost_records WHERE timestamp >= ?`,
+    [monthStart]
+  );
+  return row?.spent ?? 0;
+}
+
+export function getBudgetStatus(config: BudgetConfig): BudgetStatus {
+  const dailySpent = getDailySpend();
+  const monthlySpent = getMonthlySpend();
+  return {
+    daily: {
+      spent: dailySpent,
+      limit: config.dailyBudget,
+      remaining: Math.max(0, config.dailyBudget - dailySpent),
+      percentUsed: config.dailyBudget > 0 ? (dailySpent / config.dailyBudget) * 100 : 0,
+    },
+    monthly: {
+      spent: monthlySpent,
+      limit: config.monthlyBudget,
+      remaining: Math.max(0, config.monthlyBudget - monthlySpent),
+      percentUsed: config.monthlyBudget > 0 ? (monthlySpent / config.monthlyBudget) * 100 : 0,
+    },
+  };
+}
+
+export function getCostSummary(days = 30): CostSummary {
+  const since = new Date(Date.now() - days * 86400000).toISOString();
+
+  const total = queryOne<{
+    totalSpent: number;
+    totalRequests: number;
+    totalInputTokens: number;
+    totalOutputTokens: number;
+    cacheHits: number;
+    averageLatencyMs: number;
+  }>(
+    `SELECT
+      COALESCE(SUM(total_cost), 0) as totalSpent,
+      COUNT(*) as totalRequests,
+      COALESCE(SUM(input_tokens), 0) as totalInputTokens,
+      COALESCE(SUM(output_tokens), 0) as totalOutputTokens,
+      COALESCE(SUM(CASE WHEN cached = 1 THEN 1 ELSE 0 END), 0) as cacheHits,
+      COALESCE(AVG(latency_ms), 0) as averageLatencyMs
+    FROM cost_records
+    WHERE timestamp >= ?`,
+    [since]
+  );
+
+  const providerRows = queryAll<{
+    provider: string;
+    totalSpent: number;
+    requestCount: number;
+    inputTokens: number;
+    outputTokens: number;
+  }>(
+    `SELECT provider,
+      COALESCE(SUM(total_cost), 0) as totalSpent,
+      COUNT(*) as requestCount,
+      COALESCE(SUM(input_tokens), 0) as inputTokens,
+      COALESCE(SUM(output_tokens), 0) as outputTokens
+    FROM cost_records
+    WHERE timestamp >= ?
+    GROUP BY provider`,
+    [since]
+  );
+
+  const modelRows = queryAll<{
+    model: string;
+    totalSpent: number;
+    requestCount: number;
+    inputTokens: number;
+    outputTokens: number;
+    averageLatencyMs: number;
+  }>(
+    `SELECT model,
+      COALESCE(SUM(total_cost), 0) as totalSpent,
+      COUNT(*) as requestCount,
+      COALESCE(SUM(input_tokens), 0) as inputTokens,
+      COALESCE(SUM(output_tokens), 0) as outputTokens,
+      COALESCE(AVG(latency_ms), 0) as averageLatencyMs
+    FROM cost_records
+    WHERE timestamp >= ?
+    GROUP BY model`,
+    [since]
+  );
+
+  const totalSpent = total?.totalSpent ?? 0;
+  const totalRequests = total?.totalRequests ?? 0;
+  const totalInputTokens = total?.totalInputTokens ?? 0;
+  const totalOutputTokens = total?.totalOutputTokens ?? 0;
+  const cacheHits = total?.cacheHits ?? 0;
+  const averageLatencyMs = total?.averageLatencyMs ?? 0;
+
+  const gpt4OutputCost = (totalOutputTokens / 1000000) * 30;
+  const gpt4InputCost = (totalInputTokens / 1000000) * 10;
+  const savingsVsGpt4 = Math.max(0, (gpt4InputCost + gpt4OutputCost) - totalSpent);
+
+  const byProvider: Record<string, ProviderSummary> = {};
+  for (const row of providerRows) {
+    byProvider[row.provider] = {
+      totalSpent: row.totalSpent,
+      requestCount: row.requestCount,
+      inputTokens: row.inputTokens,
+      outputTokens: row.outputTokens,
+    };
+  }
+
+  const byModel: Record<string, ModelSummary> = {};
+  for (const row of modelRows) {
+    byModel[row.model] = {
+      totalSpent: row.totalSpent,
+      requestCount: row.requestCount,
+      inputTokens: row.inputTokens,
+      outputTokens: row.outputTokens,
+      averageLatencyMs: row.averageLatencyMs,
+    };
+  }
+
+  return {
+    totalSpent,
+    totalRequests,
+    totalInputTokens,
+    totalOutputTokens,
+    cacheHitRate: totalRequests > 0 ? cacheHits / totalRequests : 0,
+    averageLatencyMs,
+    byProvider: byProvider as CostSummary['byProvider'],
+    byModel,
+    savingsVsGpt4,
+  };
+}
+
+// ---- Cache Operations ----
+
+export function getCacheEntry(promptHash: string): CacheEntry | null {
+  const database = getDb();
+  const now = Date.now();
+
+  // Clean expired entries
+  database.run('DELETE FROM cache_entries WHERE expires_at < ?', [now]);
+
+  const row = queryOne<Record<string, unknown>>(
+    `SELECT prompt_hash, model, provider, response, input_tokens, output_tokens,
+            cost, created_at, expires_at, hit_count
+     FROM cache_entries WHERE prompt_hash = ? AND expires_at > ?`,
+    [promptHash, now]
+  );
+
+  if (row) {
+    database.run(
+      'UPDATE cache_entries SET hit_count = hit_count + 1 WHERE prompt_hash = ?',
+      [promptHash]
+    );
+    const entry: CacheEntry = {
+      promptHash: row['prompt_hash'] as string,
+      model: row['model'] as string,
+      provider: row['provider'] as ProviderId,
+      response: row['response'] as string,
+      inputTokens: row['input_tokens'] as number,
+      outputTokens: row['output_tokens'] as number,
+      cost: row['cost'] as number,
+      createdAt: row['created_at'] as number,
+      expiresAt: row['expires_at'] as number,
+      hitCount: ((row['hit_count'] as number) ?? 0) + 1,
+    };
+    return entry;
+  }
+
+  return null;
+}
+
+export function setCacheEntry(entry: CacheEntry): void {
+  const database = getDb();
+  database.run(
+    `INSERT OR REPLACE INTO cache_entries
+      (prompt_hash, model, provider, response, input_tokens, output_tokens, cost, created_at, expires_at, hit_count)
+    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+    [
+      entry.promptHash,
+      entry.model,
+      entry.provider,
+      entry.response,
+      entry.inputTokens,
+      entry.outputTokens,
+      entry.cost,
+      entry.createdAt,
+      entry.expiresAt,
+      entry.hitCount,
+    ]
+  );
+  saveDb();
+}
+
+export function clearCache(): number {
+  const database = getDb();
+  // Get count before clearing
+  const countRow = queryOne<{ cnt: number }>('SELECT COUNT(*) as cnt FROM cache_entries');
+  database.run('DELETE FROM cache_entries');
+  saveDb();
+  return countRow?.cnt ?? 0;
+}
+
+// ---- Budget Config Operations ----
+
+export function getBudgetConfig(): BudgetConfig {
+  const dailyRow = queryOne<{ value: string }>("SELECT value FROM budget_config WHERE key = 'dailyBudget'");
+  const monthlyRow = queryOne<{ value: string }>("SELECT value FROM budget_config WHERE key = 'monthlyBudget'");
+  const capRow = queryOne<{ value: string }>("SELECT value FROM budget_config WHERE key = 'perRequestCap'");
+
+  return {
+    dailyBudget: dailyRow ? parseFloat(dailyRow.value) : 10,
+    monthlyBudget: monthlyRow ? parseFloat(monthlyRow.value) : 200,
+    perRequestCap: capRow ? parseFloat(capRow.value) : 1,
+  };
+}
+
+export function setBudgetConfig(config: Partial<BudgetConfig>): void {
+  const database = getDb();
+  if (config.dailyBudget !== undefined) {
+    database.run("INSERT OR REPLACE INTO budget_config (key, value) VALUES ('dailyBudget', ?)", [
+      config.dailyBudget.toString(),
+    ]);
+  }
+  if (config.monthlyBudget !== undefined) {
+    database.run("INSERT OR REPLACE INTO budget_config (key, value) VALUES ('monthlyBudget', ?)", [
+      config.monthlyBudget.toString(),
+    ]);
+  }
+  if (config.perRequestCap !== undefined) {
+    database.run("INSERT OR REPLACE INTO budget_config (key, value) VALUES ('perRequestCap', ?)", [
+      config.perRequestCap.toString(),
+    ]);
+  }
+  saveDb();
+}
+
+// ---- Dashboard Data ----
+
+export function getDashboardData(): DashboardData {
+  const summary = getCostSummary(30);
+  const budget = getBudgetStatus(getBudgetConfig());
+
+  const recentReqRows = queryAll<Record<string, unknown>>(
+    `SELECT timestamp as time, model, provider, task_type as taskType,
+      total_cost as cost, (input_tokens + output_tokens) as tokens, cached
+    FROM cost_records
+    ORDER BY timestamp DESC
+    LIMIT 50`
+  );
+
+  const recentReqs = recentReqRows.map((r) => ({
+    time: r['time'] as string,
+    model: r['model'] as string,
+    provider: r['provider'] as string,
+    taskType: r['taskType'] as string,
+    cost: r['cost'] as number,
+    tokens: r['tokens'] as number,
+    cached: Boolean(r['cached']),
+  }));
+
+  const costByDay = queryAll<{ date: string; cost: number; requests: number }>(
+    `SELECT DATE(timestamp) as date, SUM(total_cost) as cost, COUNT(*) as requests
+    FROM cost_records
+    WHERE timestamp >= datetime('now', '-30 days')
+    GROUP BY DATE(timestamp)
+    ORDER BY date DESC`
+  );
+
+  const modelDist = queryAll<{ model: string; count: number; cost: number }>(
+    `SELECT model, COUNT(*) as count, SUM(total_cost) as cost
+    FROM cost_records
+    WHERE timestamp >= datetime('now', '-30 days')
+    GROUP BY model
+    ORDER BY count DESC`
+  );
+
+  return {
+    overview: {
+      totalSpent: summary.totalSpent,
+      totalRequests: summary.totalRequests,
+      totalSaved: summary.savingsVsGpt4,
+      cacheHitRate: summary.cacheHitRate,
+    },
+    budget,
+    recentRequests: recentReqs,
+    costByDay,
+    modelDistribution: modelDist,
+  };
+}
+
+// ---- Cleanup ----
+
+export function closeDb(): void {
+  if (saveTimer) {
+    clearInterval(saveTimer);
+    saveTimer = null;
+  }
+  if (db) {
+    saveDb();
+    db.close();
+    db = null;
+  }
+}