gbrain/test/e2e/search-quality.test.ts

/**
 * Search Quality E2E Tests
 *
 * Tests the full search pipeline against PGLite with seeded pages and
 * structured mock embeddings (basis vectors). No OpenAI API calls needed.
 *
 * Validates: compiled truth boost, detail parameter, source-aware dedup,
 * chunk_id/chunk_index in results, and getEmbeddingsByChunkIds.
 */

import { describe, test, expect, beforeAll, afterAll } from 'bun:test';
import { PGLiteEngine } from '../../src/core/pglite-engine.ts';
import type { ChunkInput, SearchResult } from '../../src/core/types.ts';

let engine: PGLiteEngine;

// Create a basis vector embedding: dimension `idx` is 1.0, rest are 0.0
function basisEmbedding(idx: number, dim = 1536): Float32Array {
  const emb = new Float32Array(dim);
  emb[idx % dim] = 1.0;
  return emb;
}

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

  // Seed test pages with compiled_truth + timeline chunks
  await engine.putPage('people/pedro', {
    type: 'person',
    title: 'Pedro Franceschi',
    compiled_truth: 'Pedro is the co-founder of Brex. Expert in fintech and payments infrastructure.',
    timeline: '2024-03-15: Met Pedro at YC dinner. Discussed AI security.',
  });

  // Seed chunks with structured embeddings
  const pedroChunks: ChunkInput[] = [
    {
      chunk_index: 0,
      chunk_text: 'Pedro is the co-founder of Brex. Expert in fintech and payments infrastructure.',
      chunk_source: 'compiled_truth',
      embedding: basisEmbedding(0), // direction 0 = fintech/compiled truth
      token_count: 15,
    },
    {
      chunk_index: 1,
      chunk_text: '2024-03-15: Met Pedro at YC dinner. Discussed AI security and Crab Trap.',
      chunk_source: 'timeline',
      embedding: basisEmbedding(1), // direction 1 = meeting/timeline
      token_count: 18,
    },
  ];
  await engine.upsertChunks('people/pedro', pedroChunks);

  await engine.putPage('companies/variant', {
    type: 'company',
    title: 'Variant Fund',
    compiled_truth: 'Variant is a crypto-native investment firm focused on web3 ownership economy.',
    timeline: '2024-06-01: Variant announced new fund.',
  });

  const variantChunks: ChunkInput[] = [
    {
      chunk_index: 0,
      chunk_text: 'Variant is a crypto-native investment firm focused on web3 ownership economy.',
      chunk_source: 'compiled_truth',
      embedding: basisEmbedding(2),
      token_count: 14,
    },
    {
      chunk_index: 1,
      chunk_text: '2024-06-01: Variant announced new fund. $450M raised.',
      chunk_source: 'timeline',
      embedding: basisEmbedding(3),
      token_count: 12,
    },
  ];
  await engine.upsertChunks('companies/variant', variantChunks);

  await engine.putPage('concepts/ai-philosophy', {
    type: 'concept',
    title: 'AI Changes Who Gets to Build',
    compiled_truth: 'AI democratizes building. The marginal cost of creation approaches zero.',
    timeline: '2024-01-10: First wrote about AI and building access.',
  });

  const aiChunks: ChunkInput[] = [
    {
      chunk_index: 0,
      chunk_text: 'AI democratizes building. The marginal cost of creation approaches zero. This changes who gets to build.',
      chunk_source: 'compiled_truth',
      embedding: basisEmbedding(4),
      token_count: 20,
    },
    {
      chunk_index: 1,
      chunk_text: '2024-01-10: First wrote about AI and building access. Shared on X.',
      chunk_source: 'timeline',
      embedding: basisEmbedding(5),
      token_count: 15,
    },
  ];
  await engine.upsertChunks('concepts/ai-philosophy', aiChunks);
}, 60_000);

afterAll(async () => {
  await engine.disconnect();
});

describe('SearchResult fields', () => {
  test('keyword search returns chunk_id and chunk_index', async () => {
    const results = await engine.searchKeyword('Pedro');
    expect(results.length).toBeGreaterThan(0);
    const r = results[0];
    expect(r.chunk_id).toBeDefined();
    expect(typeof r.chunk_id).toBe('number');
    expect(r.chunk_index).toBeDefined();
    expect(typeof r.chunk_index).toBe('number');
  });

  test('vector search returns chunk_id and chunk_index', async () => {
    const results = await engine.searchVector(basisEmbedding(0));
    expect(results.length).toBeGreaterThan(0);
    const r = results[0];
    expect(r.chunk_id).toBeDefined();
    expect(typeof r.chunk_id).toBe('number');
    expect(r.chunk_index).toBeDefined();
    expect(typeof r.chunk_index).toBe('number');
  });
});

describe('detail parameter', () => {
  test('detail=low returns only compiled_truth chunks', async () => {
    const results = await engine.searchKeyword('Pedro', { detail: 'low' });
    for (const r of results) {
      expect(r.chunk_source).toBe('compiled_truth');
    }
  });

  test('detail=high returns all chunk sources', async () => {
    const results = await engine.searchKeyword('Pedro', { detail: 'high' });
    // Should include at least compiled_truth (might include timeline depending on tsvector match)
    expect(results.length).toBeGreaterThan(0);
  });

  test('detail=low on vector search filters to compiled_truth', async () => {
    // Use a timeline-direction embedding — with detail=low, should get no results
    // or only compiled_truth results
    const results = await engine.searchVector(basisEmbedding(1), { detail: 'low' });
    for (const r of results) {
      expect(r.chunk_source).toBe('compiled_truth');
    }
  });

  test('default detail (medium) returns all sources', async () => {
    const results = await engine.searchKeyword('Pedro');
    // No filter applied, should return whatever matches
    expect(results.length).toBeGreaterThan(0);
  });
});

describe('getEmbeddingsByChunkIds', () => {
  test('returns embeddings for valid chunk IDs', async () => {
    const searchResults = await engine.searchVector(basisEmbedding(0));
    expect(searchResults.length).toBeGreaterThan(0);

    const ids = searchResults.map(r => r.chunk_id).filter((id): id is number => id != null);
    const embMap = await engine.getEmbeddingsByChunkIds(ids);

    expect(embMap.size).toBeGreaterThan(0);
    for (const [id, emb] of embMap) {
      expect(emb).toBeInstanceOf(Float32Array);
      expect(emb.length).toBe(1536);
    }
  });

  test('returns empty map for empty ID list', async () => {
    const embMap = await engine.getEmbeddingsByChunkIds([]);
    expect(embMap.size).toBe(0);
  });

  test('returns empty map for non-existent IDs', async () => {
    const embMap = await engine.getEmbeddingsByChunkIds([999999, 999998]);
    expect(embMap.size).toBe(0);
  });
});

describe('keyword search without DISTINCT ON', () => {
  test('returns multiple chunks per page', async () => {
    // Search for something that matches a page with multiple chunks
    const results = await engine.searchKeyword('Pedro', { limit: 10 });
    const pedroChunks = results.filter(r => r.slug === 'people/pedro');
    // Should be able to return more than 1 chunk per page
    // (depends on tsvector matching — Pedro is in page title/search_vector)
    expect(results.length).toBeGreaterThan(0);
  });
});

describe('compiled truth boost (vector search validates ordering)', () => {
  test('compiled_truth chunks rank first with basis vector queries', async () => {
    // Query with the compiled_truth direction for Pedro (basis 0)
    const results = await engine.searchVector(basisEmbedding(0), { limit: 5 });
    expect(results.length).toBeGreaterThan(0);
    // The closest result should be the compiled_truth chunk (basis 0)
    expect(results[0].chunk_source).toBe('compiled_truth');
    expect(results[0].slug).toBe('people/pedro');
  });

  test('timeline chunks rank first when queried with timeline direction', async () => {
    // Query with the timeline direction for Pedro (basis 1)
    const results = await engine.searchVector(basisEmbedding(1), { limit: 5 });
    expect(results.length).toBeGreaterThan(0);
    expect(results[0].chunk_source).toBe('timeline');
    expect(results[0].slug).toBe('people/pedro');
  });
});