moving

packages/agent/tests/iterative-optimizer.test.ts

@@ -0,0 +1,655 @@
+import { describe, it, expect, vi } from 'vitest';
+import {
+  IterativeGEPA,
+  paretoFront,
+  scoreCandidate,
+  aggregateScores,
+  pickWinner,
+  pickSample,
+  type Candidate,
+  type CandidateScore,
+  type IterativeGEPAOptions,
+  type MutateFn,
+} from '../src/iterative-optimizer.js';
+import type { EvalExample } from '../src/eval-dataset.js';
+import type { JudgeScore } from '../src/judge.js';
+
+/**
+ * H-09 G3 — unit tests mock the judge, so they opt into the running-judge
+ * safety-check bypass. Production code paths must still wrap with
+ * makeRunningJudge(); the check in IterativeGEPA enforces that.
+ */
+function runGEPA(options: Omit<IterativeGEPAOptions, 'allowBareJudge'>) {
+  const gepa = new IterativeGEPA();
+  return gepa.run({ ...options, allowBareJudge: true });
+}
+
+// ── Fixtures ───────────────────────────────────────────────────
+
+function makeExamples(n: number): EvalExample[] {
+  return Array.from({ length: n }, (_, i) => ({
+    input: `question ${i}`,
+    expected_output: `answer ${i}`,
+    metadata: { source: 'trace' as const },
+  }));
+}
+
+function makeScore(overall: number, extra: Partial<JudgeScore> = {}): JudgeScore {
+  return {
+    overall,
+    weighted: overall,
+    correctness: extra.correctness ?? overall,
+    procedureFollowing: extra.procedureFollowing ?? overall,
+    conciseness: extra.conciseness ?? overall,
+    lengthPenalty: extra.lengthPenalty ?? 1,
+    feedback: extra.feedback ?? `feedback for score ${overall}`,
+    parsed: true,
+  };
+}
+
+function makeCandidate(id: string, score: CandidateScore | null): Candidate {
+  return {
+    id,
+    prompt: `prompt-${id}`,
+    generation: 0,
+    parent: null,
+    strategy: 'baseline',
+    score,
+    perExample: [],
+  };
+}
+
+function makeCandidateScore(
+  overall: number,
+  dims: Partial<{ correctness: number; procedureFollowing: number; conciseness: number; lengthPenalty: number; n: number; weaknessFeedback: string[] }> = {},
+): CandidateScore {
+  return {
+    overall,
+    correctness: dims.correctness ?? overall,
+    procedureFollowing: dims.procedureFollowing ?? overall,
+    conciseness: dims.conciseness ?? overall,
+    lengthPenalty: dims.lengthPenalty ?? 1,
+    n: dims.n ?? 10,
+    weaknessFeedback: dims.weaknessFeedback ?? [],
+  };
+}
+
+// Fake judge that computes a deterministic score based on candidate text length.
+// Longer prompts → higher correctness, shorter → higher conciseness.
+function makeFakeJudge(bias: 'long' | 'short' | 'uniform' = 'uniform') {
+  return {
+    async score(args: { input: string; expected: string; actual: string }): Promise<JudgeScore> {
+      const len = args.actual.length;
+      const corr = bias === 'long' ? Math.min(1, len / 50) : bias === 'short' ? Math.max(0, 1 - len / 100) : 0.5;
+      const proc = 0.6;
+      const conc = bias === 'short' ? 0.9 : 0.5;
+      const overall = 0.5 * corr + 0.3 * proc + 0.2 * conc;
+      return {
+        overall, weighted: overall,
+        correctness: corr,
+        procedureFollowing: proc,
+        conciseness: conc,
+        lengthPenalty: 1,
+        feedback: `len=${len}`,
+        parsed: true,
+      };
+    },
+  };
+}
+
+// ── Pure helpers ───────────────────────────────────────────────
+
+describe('aggregateScores', () => {
+  it('handles empty input', () => {
+    const agg = aggregateScores([]);
+    expect(agg.n).toBe(0);
+    expect(agg.overall).toBe(0);
+    expect(agg.weaknessFeedback).toEqual([]);
+  });
+
+  it('averages component dimensions', () => {
+    const agg = aggregateScores([
+      makeScore(0.8, { correctness: 0.9, procedureFollowing: 0.7, conciseness: 0.8 }),
+      makeScore(0.6, { correctness: 0.7, procedureFollowing: 0.5, conciseness: 0.6 }),
+    ]);
+    expect(agg.n).toBe(2);
+    expect(agg.correctness).toBeCloseTo(0.8, 5);
+    expect(agg.procedureFollowing).toBeCloseTo(0.6, 5);
+    expect(agg.conciseness).toBeCloseTo(0.7, 5);
+    expect(agg.overall).toBeCloseTo(0.7, 5);
+  });
+
+  it('surfaces worst-3 feedback lines in weaknessFeedback', () => {
+    const agg = aggregateScores([
+      makeScore(0.9, { feedback: 'best' }),
+      makeScore(0.1, { feedback: 'worst' }),
+      makeScore(0.2, { feedback: 'second-worst' }),
+      makeScore(0.3, { feedback: 'third-worst' }),
+      makeScore(0.8, { feedback: 'good' }),
+    ]);
+    // Worst 3 should contain 'worst', 'second-worst', 'third-worst' in some order
+    expect(agg.weaknessFeedback).toContain('worst');
+    expect(agg.weaknessFeedback).toContain('second-worst');
+    expect(agg.weaknessFeedback).toContain('third-worst');
+    expect(agg.weaknessFeedback).not.toContain('best');
+  });
+
+  it('drops empty feedback from weakness list', () => {
+    const agg = aggregateScores([
+      makeScore(0.1, { feedback: '' }),
+      makeScore(0.2, { feedback: 'real feedback' }),
+    ]);
+    expect(agg.weaknessFeedback).toEqual(['real feedback']);
+  });
+});
+
+describe('paretoFront', () => {
+  it('keeps a single member when only one has a score', () => {
+    const a = makeCandidate('a', makeCandidateScore(0.8));
+    const b = makeCandidate('b', null);
+    expect(paretoFront([a, b])).toHaveLength(1);
+  });
+
+  it('removes strictly-dominated candidates', () => {
+    const weak = makeCandidate('weak', makeCandidateScore(0.5, {
+      correctness: 0.5, procedureFollowing: 0.5, conciseness: 0.5,
+    }));
+    const strong = makeCandidate('strong', makeCandidateScore(0.8, {
+      correctness: 0.8, procedureFollowing: 0.8, conciseness: 0.8,
+    }));
+    const front = paretoFront([weak, strong]);
+    expect(front).toHaveLength(1);
+    expect(front[0].id).toBe('strong');
+  });
+
+  it('keeps trade-off candidates on the front', () => {
+    const accurate = makeCandidate('accurate', makeCandidateScore(0.7, {
+      correctness: 0.9, procedureFollowing: 0.6, conciseness: 0.5,
+    }));
+    const concise = makeCandidate('concise', makeCandidateScore(0.7, {
+      correctness: 0.6, procedureFollowing: 0.6, conciseness: 0.9,
+    }));
+    const front = paretoFront([accurate, concise]);
+    expect(front).toHaveLength(2);
+  });
+
+  it('returns all candidates when they have identical scores', () => {
+    const a = makeCandidate('a', makeCandidateScore(0.5));
+    const b = makeCandidate('b', makeCandidateScore(0.5));
+    const front = paretoFront([a, b]);
+    expect(front).toHaveLength(2);
+  });
+
+  it('ignores candidates with null scores entirely', () => {
+    const scored = makeCandidate('scored', makeCandidateScore(0.5));
+    const unscored = makeCandidate('unscored', null);
+    const front = paretoFront([scored, unscored]);
+    expect(front).toHaveLength(1);
+    expect(front[0].id).toBe('scored');
+  });
+});
+
+describe('pickWinner', () => {
+  it('picks highest overall score', () => {
+    const a = makeCandidate('a', makeCandidateScore(0.6));
+    const b = makeCandidate('b', makeCandidateScore(0.9));
+    expect(pickWinner([a, b]).id).toBe('b');
+  });
+
+  it('breaks ties using length penalty (higher = more concise, preferred)', () => {
+    const verbose = makeCandidate('verbose', makeCandidateScore(0.7, { lengthPenalty: 0.6 }));
+    const concise = makeCandidate('concise', makeCandidateScore(0.7, { lengthPenalty: 0.9 }));
+    expect(pickWinner([verbose, concise]).id).toBe('concise');
+  });
+
+  it('breaks further ties by earliest generation', () => {
+    const early = makeCandidate('early', makeCandidateScore(0.7, { lengthPenalty: 0.8 }));
+    const late = makeCandidate('late', makeCandidateScore(0.7, { lengthPenalty: 0.8 }));
+    late.generation = 3;
+    expect(pickWinner([early, late]).id).toBe('early');
+  });
+
+  it('throws on empty input', () => {
+    expect(() => pickWinner([])).toThrow();
+  });
+});
+
+describe('pickSample', () => {
+  const rng = () => 0.5;
+
+  it('returns empty on k=0', () => {
+    expect(pickSample(makeExamples(10), 0, rng)).toEqual([]);
+  });
+
+  it('returns full set (shuffled) when k >= length', () => {
+    const sample = pickSample(makeExamples(5), 10, rng);
+    expect(sample).toHaveLength(5);
+  });
+
+  it('is deterministic with a fixed rng', () => {
+    const rngA = (() => { let s = 1; return () => { s = (s * 9301 + 49297) % 233280; return s / 233280; }; })();
+    const rngB = (() => { let s = 1; return () => { s = (s * 9301 + 49297) % 233280; return s / 233280; }; })();
+    const a = pickSample(makeExamples(20), 5, rngA).map(e => e.input);
+    const b = pickSample(makeExamples(20), 5, rngB).map(e => e.input);
+    expect(a).toEqual(b);
+  });
+});
+
+describe('scoreCandidate', () => {
+  it('scores candidate against examples and populates the score field', async () => {
+    const cand = makeCandidate('c', null);
+    const score = await scoreCandidate(cand, makeExamples(3), makeFakeJudge('uniform'));
+    expect(score.n).toBe(3);
+    expect(cand.score).not.toBeNull();
+    expect(cand.perExample).toHaveLength(3);
+  });
+
+  it('returns zero-score aggregation when examples list is empty', async () => {
+    const cand = makeCandidate('c', null);
+    const score = await scoreCandidate(cand, [], makeFakeJudge());
+    expect(score.n).toBe(0);
+    expect(score.overall).toBe(0);
+  });
+
+  it('continues iterating if one score call throws', async () => {
+    let callCount = 0;
+    const flakyJudge = {
+      async score(): Promise<JudgeScore> {
+        callCount++;
+        if (callCount === 2) throw new Error('transient');
+        return makeScore(0.5);
+      },
+    };
+    const cand = makeCandidate('c', null);
+    const score = await scoreCandidate(cand, makeExamples(3), flakyJudge);
+    expect(score.n).toBe(2); // 2 succeeded, 1 failed
+  });
+
+  it('stops early when abort signal fires', async () => {
+    const ctrl = new AbortController();
+    const judge = makeFakeJudge();
+    // Abort before any call
+    ctrl.abort();
+    const cand = makeCandidate('c', null);
+    await scoreCandidate(cand, makeExamples(10), judge, ctrl.signal);
+    expect(cand.perExample.length).toBeLessThan(10);
+  });
+
+  // ── Concurrency ────────────────────────────────────────────────
+
+  it('accepts an options object (new API) equivalently to AbortSignal (legacy)', async () => {
+    const judge = makeFakeJudge('uniform');
+    const candA = makeCandidate('a', null);
+    const candB = makeCandidate('b', null);
+
+    const ctrl = new AbortController();
+    ctrl.abort();
+
+    await scoreCandidate(candA, makeExamples(5), judge, ctrl.signal);
+    await scoreCandidate(candB, makeExamples(5), judge, { signal: ctrl.signal });
+
+    expect(candA.perExample.length).toBe(candB.perExample.length);
+  });
+
+  it('concurrency=1 produces identical aggregate to no option (sequential baseline)', async () => {
+    const examples = makeExamples(6);
+    const candSeq = makeCandidate('seq', null);
+    const candPar = makeCandidate('par', null);
+
+    const seqScore = await scoreCandidate(candSeq, examples, makeFakeJudge('uniform'));
+    const parScore = await scoreCandidate(candPar, examples, makeFakeJudge('uniform'), { concurrency: 1 });
+
+    expect(parScore.overall).toBeCloseTo(seqScore.overall);
+    expect(parScore.n).toBe(seqScore.n);
+  });
+
+  it('with concurrency > 1, runs scores in parallel (observable via in-flight counter)', async () => {
+    let inFlight = 0;
+    let peakInFlight = 0;
+    const judge = {
+      async score(): Promise<JudgeScore> {
+        inFlight++;
+        peakInFlight = Math.max(peakInFlight, inFlight);
+        // Await a microtask + a real delay so parallel workers can accumulate.
+        await new Promise(r => setTimeout(r, 5));
+        inFlight--;
+        return makeScore(0.5);
+      },
+    };
+    const cand = makeCandidate('c', null);
+    await scoreCandidate(cand, makeExamples(8), judge, { concurrency: 4 });
+
+    expect(peakInFlight).toBe(4);
+    expect(cand.perExample).toHaveLength(8);
+  });
+
+  it('with concurrency=1, exactly 1 in-flight call at a time', async () => {
+    let inFlight = 0;
+    let peakInFlight = 0;
+    const judge = {
+      async score(): Promise<JudgeScore> {
+        inFlight++;
+        peakInFlight = Math.max(peakInFlight, inFlight);
+        await new Promise(r => setTimeout(r, 3));
+        inFlight--;
+        return makeScore(0.5);
+      },
+    };
+    const cand = makeCandidate('c', null);
+    await scoreCandidate(cand, makeExamples(5), judge, { concurrency: 1 });
+    expect(peakInFlight).toBe(1);
+  });
+
+  it('concurrency is capped at examples.length (does not start idle workers)', async () => {
+    let peakInFlight = 0;
+    let inFlight = 0;
+    const judge = {
+      async score(): Promise<JudgeScore> {
+        inFlight++;
+        peakInFlight = Math.max(peakInFlight, inFlight);
+        await new Promise(r => setTimeout(r, 2));
+        inFlight--;
+        return makeScore(0.5);
+      },
+    };
+    const cand = makeCandidate('c', null);
+    await scoreCandidate(cand, makeExamples(3), judge, { concurrency: 100 });
+    expect(peakInFlight).toBeLessThanOrEqual(3);
+  });
+
+  it('parallel mode still filters out thrown-error results without corrupting the batch', async () => {
+    let call = 0;
+    const judge = {
+      async score(): Promise<JudgeScore> {
+        call++;
+        if (call % 3 === 0) throw new Error('flaky');
+        return makeScore(0.7);
+      },
+    };
+    const cand = makeCandidate('c', null);
+    const score = await scoreCandidate(cand, makeExamples(9), judge, { concurrency: 3 });
+    // 9 total, every 3rd throws → 6 succeed.
+    expect(score.n).toBe(6);
+  });
+
+  it('parallel mode respects abort signal by not dispatching further workers', async () => {
+    const ctrl = new AbortController();
+    ctrl.abort();
+    const judge = makeFakeJudge();
+    const cand = makeCandidate('c', null);
+    await scoreCandidate(cand, makeExamples(20), judge, {
+      signal: ctrl.signal,
+      concurrency: 4,
+    });
+    expect(cand.perExample.length).toBeLessThan(20);
+  });
+});
+
+// ── IterativeGEPA end-to-end with concurrency ────────────────────
+
+describe('IterativeGEPA with concurrency', () => {
+  it('threads options.concurrency into every scoreCandidate call', async () => {
+    let peakInFlight = 0;
+    let inFlight = 0;
+    const judge = {
+      async score(): Promise<JudgeScore> {
+        inFlight++;
+        peakInFlight = Math.max(peakInFlight, inFlight);
+        await new Promise(r => setTimeout(r, 2));
+        inFlight--;
+        return makeScore(0.5);
+      },
+    };
+    const mutate: MutateFn = async ({ parent, strategy }) => `${parent.prompt} :: ${strategy}`;
+
+    await runGEPA({
+      baseline: 'base',
+      examples: makeExamples(8),
+      judge,
+      mutate,
+      populationSize: 2,
+      generations: 1,
+      microScreenSize: 4,
+      miniEvalSize: 4,
+      anchorEvalSize: 4,
+      concurrency: 3,
+    });
+
+    expect(peakInFlight).toBe(3);
+  });
+});
+
+// ── H-09 G3 · running-judge guard ──────────────────────────────
+
+describe('IterativeGEPA.run · running-judge guard (H-09 G3)', () => {
+  it('throws when a bare judge is passed without allowBareJudge', async () => {
+    const bareJudge = {
+      async score(): Promise<JudgeScore> {
+        return {
+          overall: 0.5, weighted: 0.5, correctness: 0.5,
+          procedureFollowing: 0.5, conciseness: 0.5, lengthPenalty: 1,
+          feedback: '', parsed: true,
+        };
+      },
+    };
+    const mutate: MutateFn = async ({ parent }) => `${parent.prompt} mutated`;
+    await expect(
+      new IterativeGEPA().run({
+        baseline: 'seed',
+        examples: makeExamples(3),
+        judge: bareJudge,
+        mutate,
+      }),
+    ).rejects.toThrow(/not a running judge|makeRunningJudge/);
+  });
+
+  it('accepts a judge wrapped with makeRunningJudge', async () => {
+    const { makeRunningJudge } = await import('../src/evolution-llm-wiring.js');
+    const baseJudge = {
+      async score(): Promise<JudgeScore> {
+        return {
+          overall: 0.8, weighted: 0.8, correctness: 0.8,
+          procedureFollowing: 0.8, conciseness: 0.8, lengthPenalty: 1,
+          feedback: '', parsed: true,
+        };
+      },
+    };
+    const fakeLLM = { complete: async () => 'LLM output' };
+    const wrapped = makeRunningJudge(baseJudge, fakeLLM);
+    const mutate: MutateFn = async ({ parent }) => `${parent.prompt} v2`;
+    // Should not throw — the wrapped judge carries the brand.
+    const result = await new IterativeGEPA().run({
+      baseline: 'seed',
+      examples: makeExamples(3),
+      judge: wrapped,
+      mutate,
+      populationSize: 1,
+      generations: 1,
+      microScreenSize: 1, miniEvalSize: 1, anchorEvalSize: 1,
+    });
+    expect(result.winner).toBeDefined();
+  });
+});
+
+// ── End-to-end run ─────────────────────────────────────────────
+
+describe('IterativeGEPA.run', () => {
+  it('runs through all phases and produces a winner', async () => {
+    const baseline = 'short baseline prompt';
+    const examples = makeExamples(10);
+    const judge = makeFakeJudge('uniform');
+
+    // Mutate appends text — simulates generation of children
+    const mutate: MutateFn = async ({ parent, strategy }) => {
+      return `${parent.prompt} :: ${strategy}`;
+    };
+
+    const progress: string[] = [];
+    const result = await runGEPA({
+      baseline,
+      examples,
+      judge,
+      mutate,
+      populationSize: 3,
+      generations: 2,
+      microScreenSize: 5,
+      miniEvalSize: 5,
+      anchorEvalSize: 10,
+      onProgress: (e) => progress.push(`${e.phase}@g${e.generation}`),
+    });
+
+    expect(result.winner).toBeDefined();
+    expect(result.winner.score).not.toBeNull();
+    expect(result.history.length).toBeGreaterThan(1);
+    expect(progress).toContain('start@g0');
+    expect(progress.some(p => p.startsWith('anchor'))).toBe(true);
+    expect(progress.some(p => p.startsWith('done'))).toBe(true);
+  });
+
+  it('winner dominates or matches baseline on overall score', async () => {
+    const baseline = 'x'; // very short
+    const examples = makeExamples(20);
+    // Long-biased judge: longer prompts score higher
+    const judge = makeFakeJudge('long');
+
+    // Mutate doubles length each time
+    const mutate: MutateFn = async ({ parent }) => `${parent.prompt} ${parent.prompt}more`;
+
+    const result = await runGEPA({
+      baseline,
+      examples,
+      judge,
+      mutate,
+      populationSize: 3,
+      generations: 2,
+      microScreenSize: 5,
+      miniEvalSize: 5,
+      anchorEvalSize: 15,
+    });
+
+    // Winner should have evolved (not baseline prompt).
+    expect(result.winner.prompt.length).toBeGreaterThan(baseline.length);
+    expect(result.improved).toBe(true);
+    expect(result.delta).toBeGreaterThan(0);
+  });
+
+  it('passes weakness feedback from parent into mutate()', async () => {
+    const captured: Array<{ strategy: string; feedbacks: string[] }> = [];
+    const mutate: MutateFn = async ({ parent, strategy, weaknessFeedback }) => {
+      captured.push({ strategy, feedbacks: weaknessFeedback });
+      return `${parent.prompt} mutated`;
+    };
+
+    await runGEPA({
+      baseline: 'baseline',
+      examples: makeExamples(10),
+      judge: makeFakeJudge(),
+      mutate,
+      populationSize: 2,
+      generations: 1,
+      microScreenSize: 3,
+      miniEvalSize: 3,
+      anchorEvalSize: 5,
+    });
+
+    // 2 mutations were spawned — each received the parent's feedback
+    expect(captured.length).toBe(2);
+    for (const c of captured) {
+      expect(Array.isArray(c.feedbacks)).toBe(true);
+      // Feedback array may be empty on first gen if baseline had no weaknesses;
+      // we just assert the prop is present and is an array.
+    }
+  });
+
+  it('recovers when mutate throws (falls back to parent prompt)', async () => {
+    let mutations = 0;
+    const mutate: MutateFn = async ({ parent }) => {
+      mutations++;
+      if (mutations === 1) throw new Error('boom');
+      return `${parent.prompt} ok`;
+    };
+
+    const result = await runGEPA({
+      baseline: 'baseline',
+      examples: makeExamples(5),
+      judge: makeFakeJudge(),
+      mutate,
+      populationSize: 2,
+      generations: 1,
+      microScreenSize: 3,
+      miniEvalSize: 3,
+      anchorEvalSize: 3,
+    });
+
+    // Winner must be defined even though one mutation failed
+    expect(result.winner).toBeDefined();
+    expect(result.winner.score).not.toBeNull();
+  });
+
+  it('respects abort signal and exits early without throwing', async () => {
+    const ctrl = new AbortController();
+    const mutate: MutateFn = vi.fn(async () => 'should never run');
+
+    ctrl.abort(); // abort before starting generations
+    const result = await runGEPA({
+      baseline: 'baseline',
+      examples: makeExamples(5),
+      judge: makeFakeJudge(),
+      mutate,
+      populationSize: 2,
+      generations: 2,
+      microScreenSize: 2,
+      miniEvalSize: 2,
+      anchorEvalSize: 2,
+      signal: ctrl.signal,
+    });
+
+    // With abort before baseline scoring, no winner should evolve
+    expect(result.winner).toBeDefined();
+    expect(mutate).not.toHaveBeenCalled();
+  });
+
+  it('is deterministic with same seed', async () => {
+    const examples = makeExamples(20);
+    const judge = makeFakeJudge();
+    const mutate: MutateFn = async ({ parent, strategy }) => `${parent.prompt}::${strategy}`;
+
+    const run = () => runGEPA({
+      baseline: 'seed',
+      examples,
+      judge,
+      mutate,
+      populationSize: 3,
+      generations: 1,
+      microScreenSize: 5,
+      miniEvalSize: 5,
+      anchorEvalSize: 10,
+      seed: 42,
+    });
+
+    const r1 = await run();
+    const r2 = await run();
+
+    expect(r1.winner.prompt).toBe(r2.winner.prompt);
+    expect(r1.history.length).toBe(r2.history.length);
+  });
+
+  it('produces history containing baseline and all mutated children', async () => {
+    const result = await runGEPA({
+      baseline: 'base',
+      examples: makeExamples(5),
+      judge: makeFakeJudge(),
+      mutate: async ({ parent, strategy }) => `${parent.prompt}_${strategy}`,
+      populationSize: 3,
+      generations: 2,
+      microScreenSize: 3,
+      miniEvalSize: 3,
+      anchorEvalSize: 3,
+    });
+
+    expect(result.history[0].id).toBe('g0-baseline');
+    expect(result.history[0].strategy).toBe('baseline');
+    // 2 generations * 3 mutations + 1 baseline = 7 total
+    expect(result.history.length).toBe(7);
+  });
+});