brain/tools/observer-embedding-index.test.mjs

import { describe, it, expect } from 'vitest';
import { encodeBase64 } from './router-embedding.mjs';
import {
  cosineSimilarity,
  buildIndex,
  findNearestNeighbors,
  majorityOutcome,
  mapOutcomeToFamily,
} from './observer-embedding-index.mjs';

// Helpers — build a base64-encoded Float32Array embedding from a plain array.
function emb(arr) {
  return encodeBase64(new Float32Array(arr));
}

describe('cosineSimilarity', () => {
  it('returns 1 for identical unit vectors', () => {
    const a = new Float32Array([1, 0, 0, 0]);
    expect(cosineSimilarity(a, a)).toBeCloseTo(1, 6);
  });
  it('returns 0 for orthogonal vectors', () => {
    const a = new Float32Array([1, 0, 0, 0]);
    const b = new Float32Array([0, 1, 0, 0]);
    expect(cosineSimilarity(a, b)).toBeCloseTo(0, 6);
  });
  it('returns negative for opposed vectors', () => {
    const a = new Float32Array([1, 0, 0, 0]);
    const b = new Float32Array([-1, 0, 0, 0]);
    expect(cosineSimilarity(a, b)).toBeCloseTo(-1, 6);
  });
  it('handles unequal dimensions by returning 0 (guard against malformed input)', () => {
    expect(cosineSimilarity(new Float32Array([1, 0]), new Float32Array([1, 0, 0]))).toBe(0);
  });
  it('returns 0 if either side is null / empty', () => {
    expect(cosineSimilarity(null, new Float32Array([1]))).toBe(0);
    expect(cosineSimilarity(new Float32Array([1]), null)).toBe(0);
    expect(cosineSimilarity(new Float32Array([]), new Float32Array([]))).toBe(0);
  });
});

describe('mapOutcomeToFamily', () => {
  it('maps success / soft_success to "success"', () => {
    expect(mapOutcomeToFamily('success')).toBe('success');
    expect(mapOutcomeToFamily('soft_success')).toBe('success');
  });
  it('maps rework to "retry"', () => {
    expect(mapOutcomeToFamily('rework')).toBe('retry');
  });
  it('maps blocked / partial to "failure"', () => {
    expect(mapOutcomeToFamily('blocked')).toBe('failure');
    expect(mapOutcomeToFamily('partial')).toBe('failure');
  });
  it('returns null for unknown / unresolved outcomes', () => {
    expect(mapOutcomeToFamily('unknown')).toBeNull();
    expect(mapOutcomeToFamily(null)).toBeNull();
    expect(mapOutcomeToFamily('')).toBeNull();
  });
});

describe('buildIndex', () => {
  it('includes episodes with a base64 embedding AND a resolved outcome', () => {
    const eps = [
      { task_id: 'a', timestamps: { started_at: '2026-05-25T10:00:00Z' }, prompt_embedding_base64: emb([1, 0, 0, 0]), _inferredOutcome: 'success' },
      { task_id: 'b', timestamps: { started_at: '2026-05-25T11:00:00Z' }, prompt_embedding_base64: emb([0, 1, 0, 0]), _inferredOutcome: 'rework' },
    ];
    const idx = buildIndex(eps);
    expect(idx).toHaveLength(2);
    expect(idx[0].task_id).toBe('a');
    expect(idx[0].family).toBe('success');
    expect(idx[1].family).toBe('retry');
    expect(idx[0].embedding).toBeInstanceOf(Float32Array);
  });

  it('skips episodes without an embedding', () => {
    const eps = [
      { task_id: 'a', prompt_embedding_base64: null, _inferredOutcome: 'success' },
      { task_id: 'b', prompt_embedding_base64: emb([1, 0, 0, 0]), _inferredOutcome: 'success' },
    ];
    expect(buildIndex(eps)).toHaveLength(1);
  });

  it('skips episodes with unresolved outcome (unknown / null)', () => {
    const eps = [
      { task_id: 'a', prompt_embedding_base64: emb([1, 0, 0, 0]), _inferredOutcome: 'unknown' },
      { task_id: 'b', prompt_embedding_base64: emb([0, 1, 0, 0]), _inferredOutcome: 'success' },
    ];
    expect(buildIndex(eps)).toHaveLength(1);
  });

  it('skips episodes with broken / non-decodable embedding', () => {
    const eps = [
      { task_id: 'a', prompt_embedding_base64: 'not-base64!!!', _inferredOutcome: 'success' },
      { task_id: 'b', prompt_embedding_base64: emb([1, 0, 0, 0]), _inferredOutcome: 'success' },
    ];
    expect(buildIndex(eps)).toHaveLength(1);
  });
});

describe('findNearestNeighbors', () => {
  const idx = [
    { task_id: 'a', family: 'success', embedding: new Float32Array([1, 0, 0, 0]) },
    { task_id: 'b', family: 'success', embedding: new Float32Array([0.9, 0.4, 0, 0]) },
    { task_id: 'c', family: 'retry',   embedding: new Float32Array([0, 1, 0, 0]) },
    { task_id: 'd', family: 'failure', embedding: new Float32Array([0, 0, 1, 0]) },
    { task_id: 'e', family: 'success', embedding: new Float32Array([0.7, 0.7, 0, 0]) },
  ];

  it('returns top-k by cosine similarity, highest first', () => {
    const target = new Float32Array([1, 0, 0, 0]);
    const nn = findNearestNeighbors(target, idx, 3);
    expect(nn).toHaveLength(3);
    expect(nn[0].task_id).toBe('a');           // exact match
    expect(nn[0].similarity).toBeCloseTo(1, 6);
    expect(['b', 'e']).toContain(nn[1].task_id); // close to e1
    expect(nn[2].similarity).toBeLessThan(nn[1].similarity + 1e-6);
  });

  it('handles k larger than index size (returns all)', () => {
    const nn = findNearestNeighbors(new Float32Array([1, 0, 0, 0]), idx, 100);
    expect(nn.length).toBe(idx.length);
  });

  it('returns empty array if target is null / index empty', () => {
    expect(findNearestNeighbors(null, idx, 3)).toEqual([]);
    expect(findNearestNeighbors(new Float32Array([1, 0, 0, 0]), [], 3)).toEqual([]);
  });

  it('excludes a self-reference when excludeTaskId is passed', () => {
    const target = new Float32Array([1, 0, 0, 0]);
    const nn = findNearestNeighbors(target, idx, 3, { excludeTaskId: 'a' });
    expect(nn.find((n) => n.task_id === 'a')).toBeUndefined();
    expect(nn).toHaveLength(3);
  });
});

describe('majorityOutcome', () => {
  it('returns the dominant family when one wins outright', () => {
    expect(majorityOutcome([{ family: 'success' }, { family: 'success' }, { family: 'retry' }])).toBe('success');
  });

  it('returns "mixed" on a tie at the top', () => {
    expect(majorityOutcome([{ family: 'success' }, { family: 'retry' }])).toBe('mixed');
    expect(majorityOutcome([{ family: 'success' }, { family: 'retry' }, { family: 'failure' }])).toBe('mixed');
  });

  it('returns "no_neighbors" on empty input', () => {
    expect(majorityOutcome([])).toBe('no_neighbors');
    expect(majorityOutcome(null)).toBe('no_neighbors');
  });
});

// Analyzer integration covered separately in brain-retro-analyzer.test.mjs
// (similar_past_outcome_majority axis lands via analyze()).