moving

packages/agent/tests/evolve-schema.test.ts

@@ -0,0 +1,545 @@
+import { describe, it, expect } from 'vitest';
+import {
+  EvolveSchema,
+  addOutputField,
+  removeField,
+  editFieldDescription,
+  changeFieldType,
+  addConstraint,
+  removeConstraint,
+  reorderFields,
+  replaceOutputFields,
+  schemaComplexity,
+  aggregateSchemaScores,
+  paretoFrontSchema,
+  pickSchemaWinner,
+  generateStructureMutations,
+  generateOrderMutations,
+  generateRefinementMutations,
+  pickSample,
+  type Schema,
+  type SchemaField,
+  type SchemaCandidate,
+  type SchemaCandidateScore,
+  type SchemaExecuteFn,
+} from '../src/evolve-schema.js';
+import type { EvalExample } from '../src/eval-dataset.js';
+import type { JudgeScore } from '../src/judge.js';
+
+// ── Fixtures ───────────────────────────────────────────────────
+
+function makeField(name: string, partial: Partial<SchemaField> = {}): SchemaField {
+  return {
+    name,
+    type: 'string',
+    description: `description for ${name}`,
+    required: true,
+    constraints: [],
+    ...partial,
+  };
+}
+
+function makeSchema(fields: string[]): Schema {
+  return {
+    name: 'test',
+    fields: fields.map(f => makeField(f)),
+    version: 1,
+  };
+}
+
+function makeExamples(n: number): EvalExample[] {
+  return Array.from({ length: n }, (_, i) => ({
+    input: `q${i}`,
+    expected_output: `a${i}`,
+    metadata: { source: 'trace' as const },
+  }));
+}
+
+function makeJudgeScore(overall: number, feedback = 'ok'): JudgeScore {
+  return {
+    overall, weighted: overall,
+    correctness: overall, procedureFollowing: overall, conciseness: overall,
+    lengthPenalty: 1, feedback, parsed: true,
+  };
+}
+
+function makeCandidate(id: string, schema: Schema, score: SchemaCandidateScore | null): SchemaCandidate {
+  return {
+    id,
+    schema,
+    generation: 0,
+    parent: null,
+    mutation: 'baseline',
+    mutationLabel: 'baseline',
+    score,
+    perExample: [],
+  };
+}
+
+// Fake executor: returns a deterministic output string that scores higher
+// when schema has more fields (up to a cap) and scores lower when schema
+// is bloated beyond 5 fields.
+function makeFakeExecutor(bias: 'prefers-more' | 'prefers-less' | 'flat' = 'prefers-more'): SchemaExecuteFn {
+  return async ({ schema }) => ({
+    actual: schema.fields.map(f => `${f.name}=mock`).join('; '),
+    parsed: schema.fields.length > 0,
+  });
+}
+
+function makeJudge(bias: 'prefers-more' | 'prefers-less' | 'flat' = 'prefers-more') {
+  return {
+    async score(args: { input: string; expected: string; actual: string }): Promise<JudgeScore> {
+      const fieldCount = (args.actual.match(/=/g) || []).length;
+      let overall: number;
+      if (bias === 'prefers-more') {
+        // More fields = better, up to 6 fields, then flat
+        overall = Math.min(1, fieldCount / 6);
+      } else if (bias === 'prefers-less') {
+        overall = Math.max(0, 1 - fieldCount / 10);
+      } else {
+        overall = 0.5;
+      }
+      return makeJudgeScore(overall);
+    },
+  };
+}
+
+// ── Pure mutation functions ────────────────────────────────────
+
+describe('addOutputField', () => {
+  it('appends when no position given', () => {
+    const s = makeSchema(['a', 'b']);
+    const next = addOutputField(s, makeField('c'));
+    expect(next.fields.map(f => f.name)).toEqual(['a', 'b', 'c']);
+    expect(next.version).toBe(2);
+  });
+
+  it('inserts at given position', () => {
+    const s = makeSchema(['a', 'c']);
+    const next = addOutputField(s, makeField('b'), 1);
+    expect(next.fields.map(f => f.name)).toEqual(['a', 'b', 'c']);
+  });
+
+  it('does not mutate source', () => {
+    const s = makeSchema(['a']);
+    addOutputField(s, makeField('b'));
+    expect(s.fields.map(f => f.name)).toEqual(['a']);
+  });
+});
+
+describe('removeField', () => {
+  it('drops the named field', () => {
+    const s = makeSchema(['a', 'b', 'c']);
+    const next = removeField(s, 'b');
+    expect(next.fields.map(f => f.name)).toEqual(['a', 'c']);
+    expect(next.version).toBe(2);
+  });
+
+  it('is a no-op for unknown field', () => {
+    const s = makeSchema(['a', 'b']);
+    const next = removeField(s, 'z');
+    expect(next.fields.map(f => f.name)).toEqual(['a', 'b']);
+  });
+});
+
+describe('editFieldDescription', () => {
+  it('updates the description of the named field', () => {
+    const s = makeSchema(['a']);
+    const next = editFieldDescription(s, 'a', 'new');
+    expect(next.fields[0].description).toBe('new');
+  });
+});
+
+describe('changeFieldType', () => {
+  it('updates the type of the named field', () => {
+    const s = makeSchema(['a']);
+    const next = changeFieldType(s, 'a', 'number');
+    expect(next.fields[0].type).toBe('number');
+  });
+});
+
+describe('addConstraint / removeConstraint', () => {
+  it('adds a constraint', () => {
+    const s = makeSchema(['a']);
+    const next = addConstraint(s, 'a', { kind: 'maxLength', value: 10 });
+    expect(next.fields[0].constraints).toHaveLength(1);
+  });
+
+  it('removes the given constraint index', () => {
+    const s = addConstraint(makeSchema(['a']), 'a', { kind: 'maxLength', value: 10 });
+    const next = removeConstraint(s, 'a', 0);
+    expect(next.fields[0].constraints).toHaveLength(0);
+  });
+
+  it('is a no-op for unknown field when adding', () => {
+    const s = makeSchema(['a']);
+    const next = addConstraint(s, 'z', { kind: 'maxLength', value: 10 });
+    expect(next.fields[0].constraints).toHaveLength(0);
+  });
+});
+
+describe('reorderFields', () => {
+  it('reorders by name', () => {
+    const s = makeSchema(['a', 'b', 'c']);
+    const next = reorderFields(s, ['c', 'a', 'b']);
+    expect(next.fields.map(f => f.name)).toEqual(['c', 'a', 'b']);
+  });
+
+  it('appends fields omitted from newOrder', () => {
+    const s = makeSchema(['a', 'b', 'c']);
+    const next = reorderFields(s, ['b']);
+    // 'b' first, then the rest in original order
+    expect(next.fields.map(f => f.name)).toEqual(['b', 'a', 'c']);
+  });
+
+  it('ignores unknown field names in newOrder', () => {
+    const s = makeSchema(['a', 'b']);
+    const next = reorderFields(s, ['unknown', 'a', 'b']);
+    expect(next.fields.map(f => f.name)).toEqual(['a', 'b']);
+  });
+});
+
+describe('replaceOutputFields', () => {
+  it('replaces all fields', () => {
+    const s = makeSchema(['a', 'b']);
+    const next = replaceOutputFields(s, [makeField('x'), makeField('y')]);
+    expect(next.fields.map(f => f.name)).toEqual(['x', 'y']);
+  });
+});
+
+// ── Complexity + scoring ───────────────────────────────────────
+
+describe('schemaComplexity', () => {
+  it('returns 0 for empty schema', () => {
+    expect(schemaComplexity(makeSchema([]))).toBe(0);
+  });
+
+  it('scales with field count, constraint count, and description length', () => {
+    const simple = schemaComplexity(makeSchema(['a']));
+    const twoField = schemaComplexity(makeSchema(['a', 'b']));
+    expect(twoField).toBeGreaterThan(simple);
+
+    const withConstraint = schemaComplexity(
+      addConstraint(makeSchema(['a']), 'a', { kind: 'maxLength', value: 10 }),
+    );
+    expect(withConstraint).toBeGreaterThan(simple);
+  });
+});
+
+describe('aggregateSchemaScores', () => {
+  it('returns zero aggregate on empty results', () => {
+    const agg = aggregateSchemaScores([]);
+    expect(agg.n).toBe(0);
+    expect(agg.accuracy).toBe(0);
+    expect(agg.parseRate).toBe(0);
+  });
+
+  it('averages accuracy + computes parse rate', () => {
+    const results = [
+      { input: 'a', expected: 'A', actual: 'A', score: makeJudgeScore(0.8), parsed: true },
+      { input: 'b', expected: 'B', actual: 'B', score: makeJudgeScore(0.6), parsed: true },
+      { input: 'c', expected: 'C', actual: 'X', score: makeJudgeScore(0.2), parsed: false },
+    ];
+    const agg = aggregateSchemaScores(results);
+    expect(agg.n).toBe(3);
+    expect(agg.accuracy).toBeCloseTo((0.8 + 0.6 + 0.2) / 3, 5);
+    expect(agg.parseRate).toBeCloseTo(2 / 3, 5);
+  });
+
+  it('surfaces worst-3 example feedback', () => {
+    const results = [
+      { input: 'a', expected: 'A', actual: 'A', score: makeJudgeScore(0.9, 'best'), parsed: true },
+      { input: 'b', expected: 'B', actual: 'B', score: makeJudgeScore(0.1, 'worst'), parsed: true },
+      { input: 'c', expected: 'C', actual: 'C', score: makeJudgeScore(0.5, 'middle'), parsed: true },
+    ];
+    const agg = aggregateSchemaScores(results);
+    expect(agg.weaknessFeedback).toContain('worst');
+    expect(agg.weaknessFeedback).toContain('middle');
+  });
+});
+
+// ── Pareto ─────────────────────────────────────────────────────
+
+describe('paretoFrontSchema', () => {
+  const makeScore = (accuracy: number, complexity: number): SchemaCandidateScore => ({
+    accuracy, complexity, parseRate: 1, weaknessFeedback: [], n: 10,
+  });
+
+  it('removes strictly dominated candidates', () => {
+    const dominated = makeCandidate('d', makeSchema(['a']), makeScore(0.5, 5));
+    const better = makeCandidate('b', makeSchema(['a']), makeScore(0.8, 3));
+    const front = paretoFrontSchema([dominated, better]);
+    expect(front.map(c => c.id)).toEqual(['b']);
+  });
+
+  it('keeps trade-off candidates (accurate vs simple)', () => {
+    const accurate = makeCandidate('acc', makeSchema(['a', 'b']), makeScore(0.9, 10));
+    const simple = makeCandidate('sim', makeSchema(['a']), makeScore(0.7, 2));
+    const front = paretoFrontSchema([accurate, simple]);
+    expect(front).toHaveLength(2);
+  });
+
+  it('ignores unscored candidates', () => {
+    const scored = makeCandidate('s', makeSchema(['a']), makeScore(0.5, 1));
+    const unscored = makeCandidate('u', makeSchema(['a']), null);
+    const front = paretoFrontSchema([scored, unscored]);
+    expect(front).toHaveLength(1);
+  });
+});
+
+describe('pickSchemaWinner', () => {
+  const makeScore = (accuracy: number, complexity: number): SchemaCandidateScore => ({
+    accuracy, complexity, parseRate: 1, weaknessFeedback: [], n: 10,
+  });
+
+  it('picks highest accuracy', () => {
+    const lo = makeCandidate('lo', makeSchema(['a']), makeScore(0.5, 1));
+    const hi = makeCandidate('hi', makeSchema(['a']), makeScore(0.9, 10));
+    expect(pickSchemaWinner([lo, hi]).id).toBe('hi');
+  });
+
+  it('ties on accuracy → prefers lower complexity', () => {
+    const complex = makeCandidate('cx', makeSchema(['a', 'b']), makeScore(0.8, 10));
+    const simple = makeCandidate('sm', makeSchema(['a']), makeScore(0.8, 3));
+    expect(pickSchemaWinner([complex, simple]).id).toBe('sm');
+  });
+
+  it('throws on empty', () => {
+    expect(() => pickSchemaWinner([])).toThrow();
+  });
+});
+
+// ── Mutation generators ────────────────────────────────────────
+
+describe('generateStructureMutations', () => {
+  const rng = () => 0.5;
+
+  it('suggests adding reasoning if absent', () => {
+    const s = makeSchema(['answer']);
+    const muts = generateStructureMutations(s, 5, rng);
+    expect(muts.some(m => m.description.includes('reasoning'))).toBe(true);
+  });
+
+  it('does not re-suggest reasoning if already present', () => {
+    const s = makeSchema(['reasoning', 'answer']);
+    const muts = generateStructureMutations(s, 5, rng);
+    const hasReasoning = muts.some(m => m.description.includes('reasoning field'));
+    expect(hasReasoning).toBe(false);
+  });
+
+  it('suggests drop when schema is large', () => {
+    const s = makeSchema(['a', 'b', 'c', 'd', 'e']);
+    const muts = generateStructureMutations(s, 5, rng);
+    expect(muts.some(m => m.kind === 'remove_field')).toBe(true);
+  });
+
+  it('returns at most n mutations', () => {
+    const s = makeSchema(['answer']);
+    const muts = generateStructureMutations(s, 1, rng);
+    expect(muts.length).toBeLessThanOrEqual(1);
+  });
+});
+
+describe('generateOrderMutations', () => {
+  const rng = () => 0.5;
+
+  it('returns empty for 1-field schemas', () => {
+    expect(generateOrderMutations(makeSchema(['a']), 3, rng)).toEqual([]);
+  });
+
+  it('moves a reasoning field to the front', () => {
+    const s: Schema = {
+      name: 't', version: 1,
+      fields: [makeField('answer'), makeField('reasoning'), makeField('confidence')],
+    };
+    const muts = generateOrderMutations(s, 3, rng);
+    const applied = muts[0].apply(s);
+    expect(applied.fields[0].name).toBe('reasoning');
+  });
+
+  it('moves a confidence field to the end', () => {
+    const s: Schema = {
+      name: 't', version: 1,
+      fields: [makeField('confidence'), makeField('reasoning'), makeField('answer')],
+    };
+    const muts = generateOrderMutations(s, 3, rng);
+    const last = muts[0].apply(s);
+    // One of the generated mutations should put 'confidence' at the end
+    const someMovesConfidenceToEnd = muts.some(m => {
+      const applied = m.apply(s);
+      return applied.fields[applied.fields.length - 1].name === 'confidence';
+    });
+    expect(someMovesConfidenceToEnd).toBe(true);
+  });
+});
+
+describe('generateRefinementMutations', () => {
+  it('adds maxLength when feedback mentions verbosity', async () => {
+    const s = makeSchema(['reasoning', 'answer']);
+    const muts = await generateRefinementMutations(
+      s, ['too verbose, output wordy'], 5,
+    );
+    expect(muts.some(m => m.kind === 'add_constraint')).toBe(true);
+  });
+
+  it('adds minLength when feedback mentions incompleteness', async () => {
+    const s = makeSchema(['answer']);
+    const muts = await generateRefinementMutations(
+      s, ['response is too brief, missing detail'], 5,
+    );
+    expect(muts.some(m =>
+      m.description.includes('minLength'),
+    )).toBe(true);
+  });
+
+  it('edits field descriptions when feedback mentions format issues', async () => {
+    const s = makeSchema(['answer']);
+    const muts = await generateRefinementMutations(
+      s, ['wrong format, could not parse'], 5,
+    );
+    expect(muts.some(m => m.kind === 'edit_field_desc')).toBe(true);
+  });
+
+  it('falls back to a single description rewrite when no heuristic matches', async () => {
+    const s = makeSchema(['answer']);
+    const muts = await generateRefinementMutations(s, ['something unrelated'], 5);
+    expect(muts.length).toBeGreaterThan(0);
+    expect(muts[0].kind).toBe('edit_field_desc');
+  });
+
+  it('uses LLM-provided editor when supplied', async () => {
+    const s = makeSchema(['answer']);
+    const muts = await generateRefinementMutations(
+      s, ['wrong format'], 5,
+      async ({ field }) => `LLM-REWRITE of ${field.name}`,
+    );
+    const edit = muts.find(m => m.kind === 'edit_field_desc');
+    const applied = edit!.apply(s);
+    expect(applied.fields[0].description).toBe('LLM-REWRITE of answer');
+  });
+});
+
+// ── pickSample ─────────────────────────────────────────────────
+
+describe('pickSample', () => {
+  it('returns up to k items', () => {
+    expect(pickSample(makeExamples(5), 3, () => 0.5)).toHaveLength(3);
+  });
+
+  it('returns empty on k=0', () => {
+    expect(pickSample(makeExamples(5), 0, () => 0.5)).toEqual([]);
+  });
+});
+
+// ── EvolveSchema.run end-to-end ───────────────────────────────
+
+describe('EvolveSchema.run', () => {
+  it('runs through all phases and returns a winner', async () => {
+    const baseline = makeSchema(['answer']);
+    const result = await new EvolveSchema().run({
+      baseline,
+      examples: makeExamples(10),
+      execute: makeFakeExecutor(),
+      judge: makeJudge('prefers-more'),
+      populationSize: 3,
+      generations: 1,
+      evalSize: 5,
+      anchorEvalSize: 10,
+    });
+    expect(result.winner).toBeDefined();
+    expect(result.winner.score).not.toBeNull();
+    expect(result.history.length).toBeGreaterThan(1);
+  });
+
+  it('winner accuracy >= baseline accuracy when judge prefers richer schemas', async () => {
+    const baseline = makeSchema(['answer']);
+    const result = await new EvolveSchema().run({
+      baseline,
+      examples: makeExamples(10),
+      execute: makeFakeExecutor(),
+      judge: makeJudge('prefers-more'),
+      populationSize: 3,
+      generations: 2,
+      evalSize: 5,
+      anchorEvalSize: 10,
+    });
+    expect(result.winner.score!.accuracy).toBeGreaterThanOrEqual(
+      result.history[0].score!.accuracy,
+    );
+  });
+
+  it('emits progress events for each phase', async () => {
+    const phases: string[] = [];
+    await new EvolveSchema().run({
+      baseline: makeSchema(['answer']),
+      examples: makeExamples(5),
+      execute: makeFakeExecutor(),
+      judge: makeJudge(),
+      populationSize: 2,
+      generations: 1,
+      evalSize: 3,
+      anchorEvalSize: 3,
+      onProgress: (e) => phases.push(e.phase),
+    });
+    expect(phases).toContain('start');
+    expect(phases).toContain('structure');
+    expect(phases).toContain('order');
+    expect(phases).toContain('refinement');
+    expect(phases).toContain('anchor');
+    expect(phases).toContain('done');
+  });
+
+  it('is deterministic with same seed', async () => {
+    const cfg = () => ({
+      baseline: makeSchema(['answer']),
+      examples: makeExamples(10),
+      execute: makeFakeExecutor(),
+      judge: makeJudge('prefers-more'),
+      populationSize: 3,
+      generations: 1,
+      evalSize: 5,
+      anchorEvalSize: 8,
+      seed: 7,
+    });
+    const a = await new EvolveSchema().run(cfg());
+    const b = await new EvolveSchema().run(cfg());
+    expect(a.winner.schema.fields.map(f => f.name)).toEqual(
+      b.winner.schema.fields.map(f => f.name),
+    );
+  });
+
+  it('respects abort signal', async () => {
+    const ctrl = new AbortController();
+    ctrl.abort();
+    const result = await new EvolveSchema().run({
+      baseline: makeSchema(['answer']),
+      examples: makeExamples(10),
+      execute: makeFakeExecutor(),
+      judge: makeJudge(),
+      populationSize: 2,
+      generations: 2,
+      evalSize: 3,
+      anchorEvalSize: 3,
+      signal: ctrl.signal,
+    });
+    expect(result.winner).toBeDefined();
+  });
+
+  it('records all candidates in history with generation numbers', async () => {
+    const result = await new EvolveSchema().run({
+      baseline: makeSchema(['answer']),
+      examples: makeExamples(5),
+      execute: makeFakeExecutor(),
+      judge: makeJudge(),
+      populationSize: 2,
+      generations: 2,
+      mutationMix: { structure: 2, order: 0, refinement: 0 },
+      evalSize: 3,
+      anchorEvalSize: 3,
+    });
+    const gen0 = result.history.filter(c => c.generation === 0);
+    const gen1 = result.history.filter(c => c.generation === 1);
+    expect(gen0.length).toBe(1);
+    expect(gen1.length).toBeGreaterThan(0);
+  });
+});