@claude-flow/embeddings/src/__tests__/mmr.test.ts

V3 Embedding Service - OpenAI, Transformers.js, Agentic-Flow (ONNX), Mock providers with hyperbolic embeddings, normalization, and chunking

/**
 * ADR-121 Phase 10 — MMR rerank tests.
 *
 * Coverage:
 *  - Identity case (λ=1) reduces to plain top-k by relevance
 *  - Pure diversity (λ=0) ignores query relevance
 *  - Balanced (λ=0.5) demonstrably diversifies vs plain top-k on a
 *    corpus with near-duplicate clusters
 *  - k clamping (k > candidates.length, k === 0)
 *  - Order: pickOrder reflects iteration index
 *  - averagePairwiseSimilarity: MMR result is more diverse than the
 *    plain top-k from the same candidate set
 */

import { describe, it, expect } from 'vitest';
import {
  mmrRerank,
  mmrIds,
  averagePairwiseSimilarity,
  type MmrCandidate,
} from '../mmr.js';

const DIM = 8;
function vec(values: number[]): Float32Array {
  if (values.length !== DIM) throw new Error('test bug');
  return new Float32Array(values);
}

describe('mmrRerank — basic contract', () => {
  it('returns empty array when k=0', () => {
    const out = mmrRerank(
      [{ id: 'a', vector: vec([1, 0, 0, 0, 0, 0, 0, 0]) }],
      vec([1, 0, 0, 0, 0, 0, 0, 0]),
      { k: 0 },
    );
    expect(out).toEqual([]);
  });

  it('returns empty array when candidates is empty', () => {
    const out = mmrRerank([], vec([1, 0, 0, 0, 0, 0, 0, 0]), { k: 5 });
    expect(out).toEqual([]);
  });

  it('clamps k to candidates.length', () => {
    const cands: MmrCandidate[] = [
      { id: 'a', vector: vec([1, 0, 0, 0, 0, 0, 0, 0]) },
      { id: 'b', vector: vec([0, 1, 0, 0, 0, 0, 0, 0]) },
    ];
    const out = mmrRerank(cands, vec([1, 0, 0, 0, 0, 0, 0, 0]), { k: 100 });
    expect(out.length).toBe(2);
  });

  it('preserves caller-supplied payload through the rerank', () => {
    const cands: MmrCandidate[] = [
      { id: 'a', vector: vec([1, 0, 0, 0, 0, 0, 0, 0]), payload: { source: 'docs' } },
    ];
    const out = mmrRerank(cands, vec([1, 0, 0, 0, 0, 0, 0, 0]), { k: 1 });
    expect(out[0]!.payload).toEqual({ source: 'docs' });
  });

  it('pickOrder matches iteration index', () => {
    const cands: MmrCandidate[] = [
      { id: 'a', vector: vec([1, 0, 0, 0, 0, 0, 0, 0]) },
      { id: 'b', vector: vec([0, 1, 0, 0, 0, 0, 0, 0]) },
      { id: 'c', vector: vec([0, 0, 1, 0, 0, 0, 0, 0]) },
    ];
    const out = mmrRerank(cands, vec([1, 0, 0, 0, 0, 0, 0, 0]), { k: 3 });
    expect(out.map(p => p.pickOrder)).toEqual([0, 1, 2]);
  });
});

describe('mmrRerank — lambda extremes', () => {
  // Candidates with three near-duplicates of "x" and two orthogonals.
  // Plain top-k (λ=1) returns all 3 x-duplicates first.
  // Pure diversity (λ=0) starts wherever and then picks orthogonals.
  const candidates: MmrCandidate[] = [
    { id: 'x1', vector: vec([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) },
    { id: 'x2', vector: vec([0.99, 0.01, 0.01, 0.0, 0.0, 0.0, 0.0, 0.0]) },
    { id: 'x3', vector: vec([0.98, 0.02, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) },
    { id: 'y', vector: vec([0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) },
    { id: 'z', vector: vec([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) },
  ];
  const query = vec([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);

  it('λ=1 (pure relevance) returns the 3 x-duplicates first', () => {
    const out = mmrRerank(candidates, query, { k: 3, lambda: 1.0 });
    const ids = out.map(p => p.id).sort();
    expect(ids).toEqual(['x1', 'x2', 'x3']);
  });

  it('λ=0.3 (diversity-leaning) returns at least one orthogonal in top-3', () => {
    // λ=0.5 produces tied scores on the carefully-constructed
    // duplicate set (relevance ≈ redundancy = 0.99 for the
    // x-duplicates, so the tradeoff cancels). λ=0.3 weights
    // diversity more strongly and produces a deterministic spread.
    const out = mmrRerank(candidates, query, { k: 3, lambda: 0.3 });
    const ids = out.map(p => p.id);
    // Top-1 should still be x1 (highest relevance).
    expect(ids[0]).toBe('x1');
    // Top-3 should include at least one orthogonal (diversity wins
    // over redundancy at this lambda).
    expect(ids.some(id => id === 'y' || id === 'z')).toBe(true);
  });

  it('λ=0 (pure diversity) ignores relevance — still picks first by relevance, then diverges', () => {
    const out = mmrRerank(candidates, query, { k: 3, lambda: 0.0 });
    // First pick: with λ=0, score = 0 * relevance − 1 * 0 = 0 for all
    // candidates. The first non-deterministic pick lands on whoever
    // comes first in iteration; subsequent picks favor diversity.
    expect(out.length).toBe(3);
    // The result must be a valid permutation of candidates.
    const ids = new Set(out.map(p => p.id));
    expect(ids.size).toBe(3);
  });

  it('λ ∈ [0,1] is clamped (out-of-range still produces a valid result)', () => {
    const high = mmrRerank(candidates, query, { k: 3, lambda: 2.0 });
    const low = mmrRerank(candidates, query, { k: 3, lambda: -0.5 });
    expect(high.length).toBe(3);
    expect(low.length).toBe(3);
  });
});

describe('mmrRerank — diversification is measurable', () => {
  // The MMR result on a duplicate-heavy candidate set should have
  // a lower averagePairwiseSimilarity than the plain top-k from the
  // same set. That's the whole point.
  it('λ=0.3 result is more diverse than λ=1 result on duplicate-heavy candidates', () => {
    const candidates: MmrCandidate[] = [
      { id: 'a1', vector: vec([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) },
      { id: 'a2', vector: vec([0.99, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01]) },
      { id: 'a3', vector: vec([0.98, 0.0, 0.0, 0.0, 0.0, 0.0, 0.02, 0.0]) },
      { id: 'b', vector: vec([0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) },
      { id: 'c', vector: vec([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) },
      { id: 'd', vector: vec([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]) },
    ];
    const query = vec([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);

    const plain = mmrRerank(candidates, query, { k: 3, lambda: 1.0 });
    const diverse = mmrRerank(candidates, query, { k: 3, lambda: 0.3 });

    const plainAvg = averagePairwiseSimilarity(plain);
    const diverseAvg = averagePairwiseSimilarity(diverse);

    // Diverse pick should have strictly lower average pairwise sim.
    expect(diverseAvg).toBeLessThan(plainAvg);
  });

  it('redundancy field reflects max sim to picked so far', () => {
    const cands: MmrCandidate[] = [
      { id: 'a', vector: vec([1, 0, 0, 0, 0, 0, 0, 0]) },
      { id: 'b', vector: vec([0.99, 0.01, 0, 0, 0, 0, 0, 0]) },
    ];
    const out = mmrRerank(cands, vec([1, 0, 0, 0, 0, 0, 0, 0]), { k: 2, lambda: 0.5 });
    // First pick: redundancy is 0 (nothing picked yet).
    expect(out[0]!.redundancy).toBe(0);
    // Second pick: redundancy ≈ sim(first, second) — close to 1.
    expect(out[1]!.redundancy).toBeGreaterThan(0.9);
  });
});

describe('mmrIds helper', () => {
  it('returns just the picked ids in pick order', () => {
    const cands: MmrCandidate[] = [
      { id: 'a', vector: vec([1, 0, 0, 0, 0, 0, 0, 0]) },
      { id: 'b', vector: vec([0, 1, 0, 0, 0, 0, 0, 0]) },
    ];
    const ids = mmrIds(cands, vec([1, 0, 0, 0, 0, 0, 0, 0]), { k: 2 });
    expect(ids).toEqual(['a', 'b']);
  });
});

describe('averagePairwiseSimilarity', () => {
  it('returns 0 for sets smaller than 2', () => {
    expect(averagePairwiseSimilarity([])).toBe(0);
    expect(averagePairwiseSimilarity([{
      id: 'a', vector: vec([1, 0, 0, 0, 0, 0, 0, 0]),
      mmrScore: 1, relevance: 1, redundancy: 0, pickOrder: 0,
    }])).toBe(0);
  });
});