seroxdesigncleand0badd59 months agoHistory
327 lines · 8.3 KB · TypeScript
RawDownload
// CRITICALimport type { Database } from "bun:sqlite";import { openSqliteDatabase } from "../../stores/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 = openSqliteDatabase(dbPath);    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 = openSqliteDatabase(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