qmd/test/embedding-store-integration.test.ts

/**
 * embedding-store-integration.test.ts - Tests for the
 * generateEmbeddings() / EmbeddingProvider integration in store.ts.
 *
 * Uses an in-memory SQLite + a stub EmbeddingProvider to avoid loading
 * node-llama-cpp models. Verifies:
 *   - Provider's embedBatch is called when options.embedProvider is set
 *   - Model-id guard throws ModelMismatchError on mismatch
 *   - Force re-embed bypasses the guard
 *   - getDistinctEmbeddingModels reads content_vectors correctly
 */

import { describe, test, expect, beforeEach, afterEach } from "vitest";
import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
  createStore,
  generateEmbeddings,
  getDistinctEmbeddingModels,
  insertEmbedding,
  type Store,
} from "../src/store.js";
import {
  ModelMismatchError,
  type EmbeddingProvider,
  type ProviderEmbedding,
  type ProviderHealth,
} from "../src/embedding/provider.js";

// ─────────────────────────── Stub provider ───────────────────────────────────

class StubProvider implements EmbeddingProvider {
  readonly kind = "openai" as const;
  readonly modelId: string;
  readonly dim: number;
  embedBatchCalls = 0;
  embedCalls = 0;
  totalTextsEmbedded = 0;

  constructor(modelId: string, dim = 4) {
    this.modelId = modelId;
    this.dim = dim;
  }

  getModelId(): string {
    return this.modelId;
  }
  getDimensions(): number | undefined {
    return this.dim;
  }
  async healthcheck(): Promise<ProviderHealth> {
    return { ok: true, model: this.modelId, dimensions: this.dim };
  }
  async embed(text: string): Promise<ProviderEmbedding | null> {
    this.embedCalls++;
    this.totalTextsEmbedded++;
    return { embedding: this.fakeEmbed(text), model: this.modelId };
  }
  async embedBatch(texts: string[]): Promise<(ProviderEmbedding | null)[]> {
    this.embedBatchCalls++;
    this.totalTextsEmbedded += texts.length;
    return texts.map((t) => ({ embedding: this.fakeEmbed(t), model: this.modelId }));
  }
  async dispose(): Promise<void> {}

  private fakeEmbed(text: string): number[] {
    return Array.from({ length: this.dim }, (_, i) => (text.length + i) * 0.01);
  }
}

// ─────────────────────────── Test setup ──────────────────────────────────────

let workDir: string;
let store: Store;

beforeEach(() => {
  workDir = mkdtempSync(join(tmpdir(), "qmd-store-int-test-"));
  process.env.INDEX_PATH = join(workDir, "index.sqlite");
  store = createStore(process.env.INDEX_PATH);
  // Insert content + documents with the bare-minimum schema. The content
  // body needs to be non-empty so chunkDocumentByTokens emits at least one
  // chunk per doc.
  const now = "2026-04-27T00:00:00Z";
  store.db
    .prepare(`INSERT INTO content (hash, doc, created_at) VALUES (?, ?, ?)`)
    .run("hash1", "Document one body content here that is long enough to chunk.", now);
  store.db
    .prepare(`INSERT INTO content (hash, doc, created_at) VALUES (?, ?, ?)`)
    .run("hash2", "Document two body content there with different words to chunk.", now);
  store.db
    .prepare(
      `INSERT INTO documents (hash, collection, path, title, created_at, modified_at, active) VALUES (?, ?, ?, ?, ?, ?, ?)`,
    )
    .run("hash1", "test", "one.md", "One", now, now, 1);
  store.db
    .prepare(
      `INSERT INTO documents (hash, collection, path, title, created_at, modified_at, active) VALUES (?, ?, ?, ?, ?, ?, ?)`,
    )
    .run("hash2", "test", "two.md", "Two", now, now, 1);
});

afterEach(() => {
  try {
    store.close();
  } catch { /* ignore */ }
  delete process.env.INDEX_PATH;
  rmSync(workDir, { recursive: true, force: true });
});

// ─────────────────────────── getDistinctEmbeddingModels ──────────────────────

describe("getDistinctEmbeddingModels", () => {
  test("returns [] when content_vectors is empty", () => {
    expect(getDistinctEmbeddingModels(store.db)).toEqual([]);
  });

  test("returns distinct model strings", () => {
    store.ensureVecTable(4);
    insertEmbedding(store.db, "hash1", 0, 0, new Float32Array([0.1, 0.2, 0.3, 0.4]), "embeddinggemma", "2026-04-27T00:00:00Z");
    insertEmbedding(store.db, "hash2", 0, 0, new Float32Array([0.5, 0.6, 0.7, 0.8]), "embeddinggemma", "2026-04-27T00:00:00Z");
    expect(getDistinctEmbeddingModels(store.db)).toEqual(["embeddinggemma"]);
  });

  test("returns multiple distinct models when present", () => {
    store.ensureVecTable(4);
    insertEmbedding(store.db, "hash1", 0, 0, new Float32Array([0.1, 0.2, 0.3, 0.4]), "model-a", "2026-04-27T00:00:00Z");
    insertEmbedding(store.db, "hash2", 0, 0, new Float32Array([0.5, 0.6, 0.7, 0.8]), "model-b", "2026-04-27T00:00:00Z");
    const models = getDistinctEmbeddingModels(store.db).sort();
    expect(models).toEqual(["model-a", "model-b"]);
  });
});

// ─────────────────────────── generateEmbeddings + provider ───────────────────

describe("generateEmbeddings with EmbeddingProvider", () => {
  test("uses provider.embedBatch when supplied", async () => {
    const provider = new StubProvider("embeddinggemma", 4);
    const result = await generateEmbeddings(store, {
      embedProvider: provider,
      // Use small batches to keep test fast
      maxDocsPerBatch: 64,
    });
    expect(result.docsProcessed).toBe(2);
    expect(result.chunksEmbedded).toBeGreaterThan(0);
    expect(result.errors).toBe(0);
    expect(provider.embedBatchCalls + provider.embedCalls).toBeGreaterThan(0);
    expect(provider.totalTextsEmbedded).toBeGreaterThan(0);
  });

  test("model-id guard throws ModelMismatchError on mismatch", async () => {
    // Pre-populate content_vectors with a different model id
    store.ensureVecTable(4);
    insertEmbedding(
      store.db,
      "hash1",
      0,
      0,
      new Float32Array([0.1, 0.2, 0.3, 0.4]),
      "old-model",
      "2026-04-27T00:00:00Z",
    );
    const provider = new StubProvider("new-model", 4);
    await expect(
      generateEmbeddings(store, { embedProvider: provider }),
    ).rejects.toBeInstanceOf(ModelMismatchError);
  });

  test("model-id matches → proceeds", async () => {
    store.ensureVecTable(4);
    insertEmbedding(
      store.db,
      "hash1",
      0,
      0,
      new Float32Array([0.1, 0.2, 0.3, 0.4]),
      "embeddinggemma",
      "2026-04-27T00:00:00Z",
    );
    const provider = new StubProvider("embeddinggemma", 4);
    const result = await generateEmbeddings(store, { embedProvider: provider });
    // Only hash2 needs embedding (hash1 already has one)
    expect(result.docsProcessed).toBeLessThanOrEqual(2);
    expect(result.errors).toBe(0);
  });

  test("force=true bypasses model-id guard", async () => {
    store.ensureVecTable(4);
    insertEmbedding(
      store.db,
      "hash1",
      0,
      0,
      new Float32Array([0.1, 0.2, 0.3, 0.4]),
      "old-model",
      "2026-04-27T00:00:00Z",
    );
    const provider = new StubProvider("new-model", 4);
    // force=true wipes content_vectors first → guard sees empty → no throw
    const result = await generateEmbeddings(store, {
      embedProvider: provider,
      force: true,
    });
    expect(result.docsProcessed).toBe(2);
    expect(result.errors).toBe(0);
    // Now only "new-model" should be in the DB
    expect(getDistinctEmbeddingModels(store.db)).toEqual(["new-model"]);
  });

  test("empty DB → no guard issue, anything goes", async () => {
    expect(getDistinctEmbeddingModels(store.db)).toEqual([]);
    const provider = new StubProvider("anything-id", 4);
    const result = await generateEmbeddings(store, { embedProvider: provider });
    expect(result.errors).toBe(0);
  });
});