ai-lab

import Database from "better-sqlite3";
import { homedir } from "os";
import { mkdirSync } from "fs";
import { join } from "path";

export function getDataDir(configPath: string): string {
  const p = configPath.trim() || "~/ai-lab-data";
  return p.startsWith("~/") ? join(homedir(), p.slice(2)) : p;
}

const _dbCache = new Map<string, Database.Database>();

export function getDb(dataDir: string): Database.Database {
  if (_dbCache.has(dataDir)) return _dbCache.get(dataDir)!;
  mkdirSync(dataDir, { recursive: true });
  const db = new Database(join(dataDir, "ai-lab.db"));
  db.pragma("journal_mode = WAL");
  db.exec(`
    CREATE TABLE IF NOT EXISTS prompts (
      id          INTEGER PRIMARY KEY AUTOINCREMENT,
      name        TEXT    NOT NULL,
      version     INTEGER NOT NULL DEFAULT 1,
      content     TEXT    NOT NULL,
      description TEXT    NOT NULL DEFAULT '',
      tags        TEXT    NOT NULL DEFAULT '[]',
      variables   TEXT    NOT NULL DEFAULT '[]',
      createdAt   TEXT    NOT NULL,
      UNIQUE(name, version)
    );

    CREATE TABLE IF NOT EXISTS eval_datasets (
      id          INTEGER PRIMARY KEY AUTOINCREMENT,
      name        TEXT    NOT NULL UNIQUE,
      description TEXT    NOT NULL DEFAULT '',
      tags        TEXT    NOT NULL DEFAULT '[]',
      createdAt   TEXT    NOT NULL,
      updatedAt   TEXT    NOT NULL
    );

    CREATE TABLE IF NOT EXISTS eval_cases (
      id              INTEGER PRIMARY KEY AUTOINCREMENT,
      datasetId       INTEGER NOT NULL,
      input           TEXT    NOT NULL,
      expectedOutput  TEXT    NOT NULL DEFAULT '',
      tags            TEXT    NOT NULL DEFAULT '[]',
      notes           TEXT    NOT NULL DEFAULT '',
      createdAt       TEXT    NOT NULL,
      FOREIGN KEY (datasetId) REFERENCES eval_datasets(id) ON DELETE CASCADE
    );

    CREATE TABLE IF NOT EXISTS model_results (
      id          INTEGER PRIMARY KEY AUTOINCREMENT,
      caseId      INTEGER NOT NULL,
      model       TEXT    NOT NULL,
      promptName  TEXT    NOT NULL DEFAULT '',
      promptVer   INTEGER NOT NULL DEFAULT 0,
      output      TEXT    NOT NULL,
      score       REAL,
      scoreLabel  TEXT    NOT NULL DEFAULT '',
      notes       TEXT    NOT NULL DEFAULT '',
      runAt       TEXT    NOT NULL,
      FOREIGN KEY (caseId) REFERENCES eval_cases(id) ON DELETE CASCADE
    );

    CREATE INDEX IF NOT EXISTS idx_prompts_name      ON prompts(name);
    CREATE INDEX IF NOT EXISTS idx_prompts_version   ON prompts(name, version);
    CREATE INDEX IF NOT EXISTS idx_cases_dataset     ON eval_cases(datasetId);
    CREATE INDEX IF NOT EXISTS idx_results_caseId    ON model_results(caseId);
    CREATE INDEX IF NOT EXISTS idx_results_model     ON model_results(model);
  `);
  _dbCache.set(dataDir, db);
  return db;
}

export interface PromptRow {
  id: number;
  name: string;
  version: number;
  content: string;
  description: string;
  tags: string;
  variables: string;
  createdAt: string;
}

export interface DatasetRow {
  id: number;
  name: string;
  description: string;
  tags: string;
  createdAt: string;
  updatedAt: string;
}

export interface CaseRow {
  id: number;
  datasetId: number;
  input: string;
  expectedOutput: string;
  tags: string;
  notes: string;
  createdAt: string;
}

export interface ResultRow {
  id: number;
  caseId: number;
  model: string;
  promptName: string;
  promptVer: number;
  output: string;
  score: number | null;
  scoreLabel: string;
  notes: string;
  runAt: string;
}

export function savePrompt(
  db: Database.Database,
  name: string,
  content: string,
  description: string,
  tags: string[],
  variables: string[],
): PromptRow {
  const now = new Date().toISOString();
  const latest = db.prepare(
    "SELECT MAX(version) as v FROM prompts WHERE name = ?"
  ).get(name) as { v: number | null };
  const version = (latest.v ?? 0) + 1;
  db.prepare(
    `INSERT INTO prompts (name, version, content, description, tags, variables, createdAt)
     VALUES (?, ?, ?, ?, ?, ?, ?)`
  ).run(name, version, content, description, JSON.stringify(tags), JSON.stringify(variables), now);
  return db.prepare("SELECT * FROM prompts WHERE name = ? AND version = ?").get(name, version) as PromptRow;
}

export function getPrompt(
  db: Database.Database,
  name: string,
  version?: number,
): PromptRow {
  if (version !== undefined) {
    const row = db.prepare("SELECT * FROM prompts WHERE name = ? AND version = ?").get(name, version) as PromptRow | undefined;
    if (!row) throw new Error(`Prompt '${name}' version ${version} not found.`);
    return row;
  }
  const row = db.prepare(
    "SELECT * FROM prompts WHERE name = ? ORDER BY version DESC LIMIT 1"
  ).get(name) as PromptRow | undefined;
  if (!row) throw new Error(`Prompt '${name}' not found.`);
  return row;
}

export function listPromptNames(db: Database.Database): Array<{ name: string; latestVersion: number; description: string; tags: string }> {
  return db.prepare(
    `SELECT name, MAX(version) as latestVersion, description, tags
     FROM prompts GROUP BY name ORDER BY name`
  ).all() as Array<{ name: string; latestVersion: number; description: string; tags: string }>;
}

export function listPromptVersions(db: Database.Database, name: string): PromptRow[] {
  return db.prepare(
    "SELECT * FROM prompts WHERE name = ? ORDER BY version DESC"
  ).all(name) as PromptRow[];
}

export function createDataset(
  db: Database.Database,
  name: string,
  description: string,
  tags: string[],
): DatasetRow {
  const now = new Date().toISOString();
  db.prepare(
    `INSERT INTO eval_datasets (name, description, tags, createdAt, updatedAt) VALUES (?, ?, ?, ?, ?)`
  ).run(name, description, JSON.stringify(tags), now, now);
  return db.prepare("SELECT * FROM eval_datasets WHERE name = ?").get(name) as DatasetRow;
}

export function getDataset(db: Database.Database, name: string): DatasetRow {
  const row = db.prepare("SELECT * FROM eval_datasets WHERE name = ?").get(name) as DatasetRow | undefined;
  if (!row) throw new Error(`Dataset '${name}' not found.`);
  return row;
}

export function listDatasets(db: Database.Database): Array<DatasetRow & { caseCount: number }> {
  return db.prepare(
    `SELECT d.*, COUNT(c.id) as caseCount
     FROM eval_datasets d
     LEFT JOIN eval_cases c ON c.datasetId = d.id
     GROUP BY d.id ORDER BY d.name`
  ).all() as Array<DatasetRow & { caseCount: number }>;
}

export function addEvalCase(
  db: Database.Database,
  datasetId: number,
  input: string,
  expectedOutput: string,
  tags: string[],
  notes: string,
): CaseRow {
  const now = new Date().toISOString();
  const r = db.prepare(
    `INSERT INTO eval_cases (datasetId, input, expectedOutput, tags, notes, createdAt)
     VALUES (?, ?, ?, ?, ?, ?)`
  ).run(datasetId, input, expectedOutput, JSON.stringify(tags), notes, now);
  return db.prepare("SELECT * FROM eval_cases WHERE id = ?").get(r.lastInsertRowid) as CaseRow;
}

export function getDatasetCases(db: Database.Database, datasetId: number): CaseRow[] {
  return db.prepare("SELECT * FROM eval_cases WHERE datasetId = ? ORDER BY createdAt").all(datasetId) as CaseRow[];
}

export function logModelResult(
  db: Database.Database,
  caseId: number,
  model: string,
  promptName: string,
  promptVer: number,
  output: string,
  score: number | null,
  scoreLabel: string,
  notes: string,
): ResultRow {
  const now = new Date().toISOString();
  const r = db.prepare(
    `INSERT INTO model_results (caseId, model, promptName, promptVer, output, score, scoreLabel, notes, runAt)
     VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`
  ).run(caseId, model, promptName, promptVer, output, score, scoreLabel, notes, now);
  return db.prepare("SELECT * FROM model_results WHERE id = ?").get(r.lastInsertRowid) as ResultRow;
}

export function getResultsForDataset(
  db: Database.Database,
  datasetId: number,
): Array<CaseRow & { results: ResultRow[] }> {
  const cases = getDatasetCases(db, datasetId);
  return cases.map((c) => ({
    ...c,
    results: db.prepare(
      "SELECT * FROM model_results WHERE caseId = ? ORDER BY runAt"
    ).all(c.id) as ResultRow[],
  }));
}

export function getModelStats(
  db: Database.Database,
  datasetId: number,
): Array<{ model: string; promptName: string; promptVer: number; caseCount: number; avgScore: number | null; scoredCount: number }> {
  return db.prepare(
    `SELECT mr.model, mr.promptName, mr.promptVer,
            COUNT(mr.id) as caseCount,
            AVG(mr.score) as avgScore,
            COUNT(mr.score) as scoredCount
     FROM model_results mr
     JOIN eval_cases ec ON ec.id = mr.caseId
     WHERE ec.datasetId = ?
     GROUP BY mr.model, mr.promptName, mr.promptVer
     ORDER BY avgScore DESC NULLS LAST`
  ).all(datasetId) as Array<{ model: string; promptName: string; promptVer: number; caseCount: number; avgScore: number | null; scoredCount: number }>;
}