gbrain/test/e2e/embedding-column-pglite.test.ts

/**
 * v0.36 E2E — dynamic embedding column selection (PGLite).
 *
 * Covers (per D4 + D9 + D11 + D12 + CDX-2 + CDX-3 + CDX-7 + CDX-8 + CDX-10):
 *   - Multi-column search: same query against `embedding` and against an
 *     ad-hoc `embedding_voyage` column produces different orderings
 *     consistent with the seeded vectors.
 *   - Halfvec column: ALTER TABLE ADD `embedding_ze halfvec(2560)` and
 *     confirm the `$1::halfvec(2560)` cast works.
 *   - Image branch unaffected: `embedding_image` still works via the
 *     existing operations.ts path.
 *   - cosineReScore reads from the active column, not the default
 *     (D9 — pre-fix, rescore against Voyage HNSW used OpenAI vectors).
 *   - Unknown column at hybridSearch entry throws loud.
 *   - Mid-session column switch invalidates the cache (knobs_hash v=3).
 *
 * No DATABASE_URL needed — PGLite in-memory.
 */

import { describe, test, expect, beforeAll, afterAll } from 'bun:test';
import { PGLiteEngine } from '../../src/core/pglite-engine.ts';
import { hybridSearch } from '../../src/core/search/hybrid.ts';
import {
  buildVectorCastFragment,
  EmbeddingColumnNotRegisteredError,
} from '../../src/core/search/embedding-column.ts';
import {
  configureGateway,
  resetGateway,
  __setEmbedTransportForTests,
} from '../../src/core/ai/gateway.ts';
import type { ResolvedColumn } from '../../src/core/types.ts';

let engine: PGLiteEngine;
let chunkIdA: number;
let chunkIdB: number;

const VEC1536_A = new Array(1536).fill(0).map((_, i) => 0.001 * (i % 10));
const VEC1536_B = new Array(1536).fill(0).map((_, i) => 0.002 * (i % 10));
const VEC1024_A = new Array(1024).fill(0).map((_, i) => 0.5 - 0.001 * (i % 10));
const VEC1024_B = new Array(1024).fill(0).map((_, i) => 0.4 + 0.001 * (i % 10));

beforeAll(async () => {
  engine = new PGLiteEngine();
  await engine.connect({});
  await engine.initSchema();

  // Add the ad-hoc Voyage + ZE columns the way a user with a multi-provider
  // brain has done it (outside the committed schema, per-instance ALTER).
  await (engine as any).db.exec(
    `ALTER TABLE content_chunks ADD COLUMN IF NOT EXISTS embedding_voyage vector(1024)`,
  );
  await (engine as any).db.exec(
    `ALTER TABLE content_chunks ADD COLUMN IF NOT EXISTS embedding_ze halfvec(2560)`,
  );

  // Two pages with one chunk each.
  await engine.putPage('docs/page-a', {
    type: 'concept',
    title: 'Page A — about cats',
    compiled_truth: 'Page A discusses cats and their behavior.',
  });
  await engine.putPage('docs/page-b', {
    type: 'concept',
    title: 'Page B — about dogs',
    compiled_truth: 'Page B discusses dogs and their habits.',
  });

  await engine.upsertChunks('docs/page-a', [
    { chunk_index: 0, chunk_text: 'cats behavior chunk A', chunk_source: 'compiled_truth' },
  ]);
  await engine.upsertChunks('docs/page-b', [
    { chunk_index: 0, chunk_text: 'dogs habits chunk B', chunk_source: 'compiled_truth' },
  ]);

  // Look up chunk ids.
  const rows = await engine.executeRaw<{ id: number; slug: string }>(
    `SELECT cc.id, p.slug FROM content_chunks cc JOIN pages p ON p.id = cc.page_id ORDER BY p.slug`,
  );
  chunkIdA = rows.find(r => r.slug === 'docs/page-a')!.id;
  chunkIdB = rows.find(r => r.slug === 'docs/page-b')!.id;

  // Seed vectors. Vectors are intentionally distinct between columns so
  // search orderings depend on which column the engine actually reads.
  const vecLit = (arr: number[]) => `[${arr.join(',')}]`;
  await (engine as any).db.query(
    `UPDATE content_chunks SET embedding = $1::vector WHERE id = $2`,
    [vecLit(VEC1536_A), chunkIdA],
  );
  await (engine as any).db.query(
    `UPDATE content_chunks SET embedding = $1::vector WHERE id = $2`,
    [vecLit(VEC1536_B), chunkIdB],
  );
  await (engine as any).db.query(
    `UPDATE content_chunks SET embedding_voyage = $1::vector WHERE id = $2`,
    [vecLit(VEC1024_A), chunkIdA],
  );
  await (engine as any).db.query(
    `UPDATE content_chunks SET embedding_voyage = $1::vector WHERE id = $2`,
    [vecLit(VEC1024_B), chunkIdB],
  );
});

afterAll(async () => {
  if (engine) await engine.disconnect();
  __setEmbedTransportForTests(null);
  resetGateway();
});

describe('PGLite engine: searchVector accepts ResolvedColumn descriptor (D11)', () => {
  test('vector cast routes to correct column when descriptor names embedding_voyage', async () => {
    const queryVec = new Float32Array(VEC1024_A);
    const descriptor: ResolvedColumn = {
      name: 'embedding_voyage',
      type: 'vector',
      dimensions: 1024,
      embeddingModel: 'voyage:voyage-3-large',
    };
    const results = await engine.searchVector(queryVec, {
      embeddingColumn: descriptor,
      limit: 5,
    });
    // Both pages have voyage embeddings; cosine to VEC1024_A is closer to
    // page-a (identical) than page-b. Verify ordering.
    expect(results.length).toBeGreaterThanOrEqual(1);
    expect(results[0].slug).toBe('docs/page-a');
  });

  test('halfvec cast accepted: ALTER TABLE column + $1::halfvec(N)', async () => {
    // Seed halfvec values via direct cast.
    const ze1 = `[${new Array(2560).fill(0.5).join(',')}]`;
    const ze2 = `[${new Array(2560).fill(0.6).join(',')}]`;
    await (engine as any).db.query(
      `UPDATE content_chunks SET embedding_ze = $1::halfvec WHERE id = $2`,
      [ze1, chunkIdA],
    );
    await (engine as any).db.query(
      `UPDATE content_chunks SET embedding_ze = $1::halfvec WHERE id = $2`,
      [ze2, chunkIdB],
    );

    const queryVec = new Float32Array(2560).fill(0.5);
    const descriptor: ResolvedColumn = {
      name: 'embedding_ze',
      type: 'halfvec',
      dimensions: 2560,
      embeddingModel: 'zeroentropyai:zembed-1',
    };
    const results = await engine.searchVector(queryVec, {
      embeddingColumn: descriptor,
      limit: 5,
    });
    expect(results.length).toBeGreaterThanOrEqual(1);
    // Page A's halfvec is closer to the all-0.5 query.
    expect(results[0].slug).toBe('docs/page-a');
  });

  test('legacy embedding_image literal still routes correctly', async () => {
    // We never seeded embedding_image so we expect zero results, but the
    // query MUST NOT throw — the legacy-literal path must still work
    // (no regression on the existing image branch).
    const v = new Float32Array(1024).fill(0.1);
    const results = await engine.searchVector(v, {
      embeddingColumn: 'embedding_image',
      limit: 5,
    });
    expect(Array.isArray(results)).toBe(true);
  });
});

describe('PGLite engine: getEmbeddingsByChunkIds column param (D9)', () => {
  test('default fetches from embedding (back-compat)', async () => {
    const map = await engine.getEmbeddingsByChunkIds([chunkIdA, chunkIdB]);
    expect(map.get(chunkIdA)!.length).toBe(1536);
  });

  test('column="embedding_voyage" fetches from voyage column', async () => {
    const map = await engine.getEmbeddingsByChunkIds([chunkIdA, chunkIdB], 'embedding_voyage');
    expect(map.get(chunkIdA)!.length).toBe(1024);
  });

  test('invalid column rejected at engine layer (regex guard)', async () => {
    let threw: Error | null = null;
    try {
      await engine.getEmbeddingsByChunkIds([chunkIdA], 'embed-bad-name');
    } catch (e) {
      threw = e as Error;
    }
    expect(threw).toBeInstanceOf(EmbeddingColumnNotRegisteredError);
  });
});

describe('hybridSearch + resolver — unknown column at entry (D11)', () => {
  test('unknown name in opts.embeddingColumn throws via resolver', async () => {
    // configureGateway with a transport stub so we don't hit a real API.
    configureGateway({
      embedding_model: 'openai:text-embedding-3-large',
      embedding_dimensions: 1536,
      env: { OPENAI_API_KEY: 'sk-test' },
    });
    __setEmbedTransportForTests(async () => ({
      embeddings: [new Array(1536).fill(0)],
      usage: { tokens: 0 },
    } as any));

    let threw: Error | null = null;
    try {
      await hybridSearch(engine, 'cats', {
        embeddingColumn: 'nonexistent_column',
        limit: 5,
      });
    } catch (e) {
      threw = e as Error;
    }
    expect(threw).toBeInstanceOf(EmbeddingColumnNotRegisteredError);
  });
});

describe('upsertChunks — model provenance uses gateway-resolved model, not compiled default', () => {
  // Regression (zbrain-rfi): when a caller builds ChunkInputs without an
  // explicit `model` (as src/commands/embed.ts does), the engine used to
  // stamp the compile-time DEFAULT_EMBEDDING_MODEL ('zeroentropyai:zembed-1')
  // onto content_chunks.model — even though the vector was produced by the
  // config-resolved model. That corrupted provenance the signature-drift +
  // dim-migration logic trusts. The engine must fall back to the model the
  // gateway ACTUALLY resolves at write time.
  test('unspecified chunk.model records the resolved model, not zeroentropyai:zembed-1', async () => {
    configureGateway({
      embedding_model: 'openai:text-embedding-3-large',
      embedding_dimensions: 1536,
      env: { OPENAI_API_KEY: 'sk-test' },
    });

    await engine.putPage('docs/provenance-page', {
      type: 'concept',
      title: 'Provenance test page',
      compiled_truth: 'Chunk whose model column must reflect the resolved model.',
    });
    // No `model` field on the input — the write-side fallback must fill it.
    await engine.upsertChunks('docs/provenance-page', [
      { chunk_index: 0, chunk_text: 'provenance chunk', chunk_source: 'compiled_truth' },
    ]);

    const rows = await engine.executeRaw<{ model: string }>(
      `SELECT cc.model FROM content_chunks cc
         JOIN pages p ON p.id = cc.page_id
        WHERE p.slug = 'docs/provenance-page'`,
    );
    expect(rows.length).toBe(1);
    expect(rows[0].model).toBe('openai:text-embedding-3-large');
    expect(rows[0].model).not.toBe('zeroentropyai:zembed-1');

    resetGateway();
  });

  // #3461: getEmbeddingModel() THROWS when the gateway is unconfigured — it
  // never returns falsy — so the reland's `|| resolvedModel` guard was dead
  // code and the catch path still stamped the compile-time default onto rows
  // whose vectors came from the config-resolved provider. The engine must
  // fall back to the brain's own `config.embedding_model` row instead.
  test('#3461: unconfigured gateway falls back to the brain config model, never the compiled default', async () => {
    await engine.setConfig('embedding_model', 'voyage:voyage-3-large');
    // The preload's beforeEach re-configures the gateway before every test,
    // so the reset must happen INSIDE the test body.
    resetGateway();

    await engine.putPage('docs/provenance-throw-path', {
      type: 'concept',
      title: 'Provenance throw-path page',
      compiled_truth: 'Chunk written while the gateway is unconfigured.',
    });
    await engine.upsertChunks('docs/provenance-throw-path', [
      { chunk_index: 0, chunk_text: 'throw-path provenance chunk', chunk_source: 'compiled_truth' },
    ]);

    const rows = await engine.executeRaw<{ model: string }>(
      `SELECT cc.model FROM content_chunks cc
         JOIN pages p ON p.id = cc.page_id
        WHERE p.slug = 'docs/provenance-throw-path'`,
    );
    expect(rows.length).toBe(1);
    expect(rows[0].model).toBe('voyage:voyage-3-large');

    // Restore the value initSchema wrote for the rest of the file.
    await engine.setConfig('embedding_model', 'openai:text-embedding-3-large');
  });

  // #3461 sibling: on a partial re-upsert that carries NO new embedding (the
  // exact shape `embed --stale` produces for a page's non-stale chunks), the
  // preserved vector must KEEP its original model label. The old
  // COALESCE(EXCLUDED.model, …) relabeled it with the current gateway model.
  test('#3461: preserved vector keeps its original model label on a no-embedding re-upsert', async () => {
    configureGateway({
      embedding_model: 'openai:text-embedding-3-large',
      embedding_dimensions: 1536,
      env: { OPENAI_API_KEY: 'sk-test' },
    });

    await engine.putPage('docs/provenance-preserve', {
      type: 'concept',
      title: 'Provenance preserve page',
      compiled_truth: 'Chunk embedded under model A, re-upserted under model B.',
    });
    await engine.upsertChunks('docs/provenance-preserve', [
      {
        chunk_index: 0,
        chunk_text: 'stable chunk text',
        chunk_source: 'compiled_truth',
        embedding: new Float32Array(VEC1536_A),
      },
    ]);

    // Model swap: the gateway now resolves a different model, and the
    // re-upsert (same chunk_text) carries no new embedding.
    configureGateway({
      embedding_model: 'voyage:voyage-3-large',
      embedding_dimensions: 1536,
      env: { VOYAGE_API_KEY: 'test' },
    });
    await engine.upsertChunks('docs/provenance-preserve', [
      { chunk_index: 0, chunk_text: 'stable chunk text', chunk_source: 'compiled_truth' },
    ]);

    const rows = await engine.executeRaw<{ model: string; has_embedding: boolean }>(
      `SELECT cc.model, cc.embedding IS NOT NULL AS has_embedding
         FROM content_chunks cc
         JOIN pages p ON p.id = cc.page_id
        WHERE p.slug = 'docs/provenance-preserve'`,
    );
    expect(rows.length).toBe(1);
    expect(rows[0].has_embedding).toBe(true); // vector preserved…
    expect(rows[0].model).toBe('openai:text-embedding-3-large'); // …and its label still describes it

    resetGateway();
  });
});

describe('buildVectorCastFragment — engine SQL composer (D3)', () => {
  test('vector descriptor emits $1::vector', () => {
    const r: ResolvedColumn = {
      name: 'embedding',
      type: 'vector',
      dimensions: 1536,
      embeddingModel: '',
    };
    const { col, castSql } = buildVectorCastFragment(r);
    expect(col).toBe('"embedding"');
    expect(castSql).toBe('$1::vector');
  });

  test('halfvec descriptor emits $1::halfvec(N) with parenthesized N', () => {
    const r: ResolvedColumn = {
      name: 'embedding_ze',
      type: 'halfvec',
      dimensions: 2560,
      embeddingModel: 'zeroentropyai:zembed-1',
    };
    const { col, castSql } = buildVectorCastFragment(r);
    expect(col).toBe('"embedding_ze"');
    expect(castSql).toBe('$1::halfvec(2560)');
  });
});