289 lines · 7.9 KB · TypeScript
RawDownload
// CRITICALimport { Database } from "bun:sqlite"; /** * SQLite-backed storage for peak metrics per model. */export class PeakMetricsStore {  private readonly db: Database;   /**   * Create a peak metrics store.   * @param dbPath - SQLite database path.   */  public constructor(dbPath: string) {    this.db = new Database(dbPath);    this.db.run("PRAGMA busy_timeout = 5000");    this.migrate();  }   /**   * Initialize schema.   * @returns void   */  private migrate(): void {    this.db.run(`      CREATE TABLE IF NOT EXISTS peak_metrics (        model_id TEXT PRIMARY KEY,        prefill_tps REAL,        generation_tps REAL,        ttft_ms REAL,        total_tokens INTEGER DEFAULT 0,        total_requests INTEGER DEFAULT 0,        updated_at TEXT DEFAULT CURRENT_TIMESTAMP      )    `);  }   /**   * Get peak metrics for a model.   * @param modelId - Model id.   * @returns Metrics row or null.   */  public get(modelId: string): Record<string, unknown> | null {    const row = this.db.query("SELECT * FROM peak_metrics WHERE model_id = ?").get(modelId) as Record<string, unknown> | null;    return row ? { ...row } : null;  }   /**   * Update metrics if new values are better.   * @param modelId - Model id.   * @param prefillTps - Prefill tokens per second.   * @param generationTps - Generation tokens per second.   * @param ttftMs - Time-to-first-token in ms.   * @returns Updated metrics.   */  public updateIfBetter(    modelId: string,    prefillTps?: number,    generationTps?: number,    ttftMs?: number,  ): Record<string, unknown> {    const current = this.get(modelId);    const updates: Record<string, number> = {};     if (current) {      if (prefillTps !== undefined && (current["prefill_tps"] === null || Number(prefillTps) > Number(current["prefill_tps"]))) {        updates["prefill_tps"] = prefillTps;      }      if (generationTps !== undefined && (current["generation_tps"] === null || Number(generationTps) > Number(current["generation_tps"]))) {        updates["generation_tps"] = generationTps;      }      if (ttftMs !== undefined && (current["ttft_ms"] === null || Number(ttftMs) < Number(current["ttft_ms"]))) {        updates["ttft_ms"] = ttftMs;      }    } else {      if (prefillTps !== undefined) {        updates["prefill_tps"] = prefillTps;      }      if (generationTps !== undefined) {        updates["generation_tps"] = generationTps;      }      if (ttftMs !== undefined) {        updates["ttft_ms"] = ttftMs;      }    }     if (Object.keys(updates).length > 0) {      if (current) {        const setClause = Object.keys(updates).map((key) => `${key} = ?`).join(", ");        this.db.query(`UPDATE peak_metrics SET ${setClause}, updated_at = CURRENT_TIMESTAMP WHERE model_id = ?`)          .run(...Object.values(updates), modelId);      } else {        this.db.query(`          INSERT INTO peak_metrics (model_id, prefill_tps, generation_tps, ttft_ms)          VALUES (?, ?, ?, ?)        `).run(          modelId,          updates["prefill_tps"] ?? null,          updates["generation_tps"] ?? null,          updates["ttft_ms"] ?? null,        );      }    }     return this.get(modelId) ?? {};  }   /**   * Add cumulative token and request counts.   * @param modelId - Model id.   * @param tokens - Tokens count.   * @param requests - Request count.   * @returns void   */  public addTokens(modelId: string, tokens: number, requests = 1): void {    this.db.query(`      INSERT INTO peak_metrics (model_id, total_tokens, total_requests)      VALUES (?, ?, ?)      ON CONFLICT(model_id) DO UPDATE SET        total_tokens = total_tokens + excluded.total_tokens,        total_requests = total_requests + excluded.total_requests,        updated_at = CURRENT_TIMESTAMP    `).run(modelId, tokens, requests);  }   /**   * Get all peak metrics.   * @returns List of metrics rows.   */  public getAll(): Array<Record<string, unknown>> {    const rows = this.db.query("SELECT * FROM peak_metrics ORDER BY model_id").all() as Array<Record<string, unknown>>;    return rows.map((row) => ({ ...row }));  }} /** * SQLite-backed storage for lifetime metrics. */export class LifetimeMetricsStore {  private readonly db: Database;   /**   * Create a lifetime metrics store.   * @param dbPath - SQLite database path.   */  public constructor(dbPath: string) {    this.db = new Database(dbPath);    this.migrate();  }   /**   * Initialize schema and defaults.   * @returns void   */  private migrate(): void {    this.db.run(`      CREATE TABLE IF NOT EXISTS lifetime_metrics (        key TEXT PRIMARY KEY,        value REAL NOT NULL DEFAULT 0,        updated_at TEXT DEFAULT CURRENT_TIMESTAMP      )    `);    const defaults: Array<[string, number]> = [      ["tokens_total", 0],      ["prompt_tokens_total", 0],      ["completion_tokens_total", 0],      ["energy_wh", 0],      ["uptime_seconds", 0],      ["requests_total", 0],      ["first_started_at", 0],    ];    for (const [key, value] of defaults) {      this.db.query("INSERT OR IGNORE INTO lifetime_metrics (key, value) VALUES (?, ?)").run(key, value);    }  }   /**   * Get a lifetime metric value.   * @param key - Metric key.   * @returns Metric value.   */  public get(key: string): number {    const row = this.db.query("SELECT value FROM lifetime_metrics WHERE key = ?").get(key) as { value?: number } | null;    return row?.value ?? 0;  }   /**   * Get all lifetime metrics.   * @returns Map of metric values.   */  public getAll(): Record<string, number> {    const rows = this.db.query("SELECT key, value FROM lifetime_metrics").all() as Array<{ key: string; value: number }>;    return Object.fromEntries(rows.map((row) => [row.key, row.value]));  }   /**   * Set a lifetime metric.   * @param key - Metric key.   * @param value - Metric value.   * @returns void   */  public set(key: string, value: number): void {    this.db.query(      `INSERT INTO lifetime_metrics (key, value, updated_at)       VALUES (?, ?, CURRENT_TIMESTAMP)       ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = CURRENT_TIMESTAMP`,    ).run(key, value);  }   /**   * Increment a lifetime metric.   * @param key - Metric key.   * @param delta - Increment value.   * @returns Updated value.   */  public increment(key: string, delta: number): number {    this.db.query(      `INSERT INTO lifetime_metrics (key, value, updated_at)       VALUES (?, ?, CURRENT_TIMESTAMP)       ON CONFLICT(key) DO UPDATE SET value = value + excluded.value, updated_at = CURRENT_TIMESTAMP`,    ).run(key, delta);    return this.get(key);  }   /**   * Ensure first_started_at is set.   * @returns void   */  public ensureFirstStarted(): void {    const current = this.get("first_started_at");    if (current === 0) {      this.set("first_started_at", Date.now() / 1000);    }  }   /**   * Add energy consumption in watt-hours.   * @param wattHours - Watt hours to add.   * @returns void   */  public addEnergy(wattHours: number): void {    this.increment("energy_wh", wattHours);  }   /**   * Add total tokens.   * @param tokens - Tokens to add.   * @returns void   */  public addTokens(tokens: number): void {    this.increment("tokens_total", tokens);  }   /**   * Add prompt tokens.   * @param tokens - Tokens to add.   * @returns void   */  public addPromptTokens(tokens: number): void {    this.increment("prompt_tokens_total", tokens);  }   /**   * Add completion tokens.   * @param tokens - Tokens to add.   * @returns void   */  public addCompletionTokens(tokens: number): void {    this.increment("completion_tokens_total", tokens);  }   /**   * Add uptime in seconds.   * @param seconds - Seconds to add.   * @returns void   */  public addUptime(seconds: number): void {    this.increment("uptime_seconds", seconds);  }   /**   * Add request count.   * @param count - Requests to add.   * @returns void   */  public addRequests(count = 1): void {    this.increment("requests_total", count);  }}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t