fix(ace): Ensure only input fields are passed to curator (#456)

The AxACE optimizer was incorrectly passing the entire example, including output fields, to the curator's `question_context`. This caused the curator to generate a playbook that confused inputs and outputs.

This commit fixes the issue by filtering the example object based on the program's input signature before serializing it to `question_context`.

Additionally, this commit:
- Adds a unit test to verify that the `question_context` is correctly formed.
- Removes several `console.log` statements from the test suite.
- Refactors the test suite to use the `flow()` factory function instead of the deprecated `new AxFlow()` constructor.

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: google-labs-jules[bot]

src/ax/dsp/optimizer.test.ts

@@ -1,24 +1,66 @@
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
 
 import type { AxAIService } from '../ai/types.js';
 
 import type { AxOptimizer } from './optimizer.js';
 import { AxBootstrapFewShot } from './optimizers/bootstrapFewshot.js';
+import { AxACE } from './optimizers/ace.js';
 import { AxMiPRO } from './optimizers/miproV2.js';
+import { f } from './sig.js';
+import { ax } from './template.js';
 
 // Mock dependencies
 const mockAI = {
   name: 'mock',
-  chat: async () => ({ results: [{ index: 0, content: 'mock response' }] }),
+  chat: async () => ({
+    results: [
+      {
+        index: 0,
+        content: JSON.stringify({
+          answer: 'mock student response',
+        }),
+      },
+    ],
+  }),
+  getOptions: () => ({ tracer: undefined }),
+  getLogger: () => undefined,
 } as unknown as AxAIService;
 
-// Removed unused mockProgram
-
 const mockExamples = [
   { input: 'test input', output: 'test output' },
   { input: 'test input 2', output: 'test output 2' },
 ];
 
+describe('AxACE', () => {
+  it('should call runCurator with the correct context', async () => {
+    const ace = new AxACE({
+      studentAI: mockAI,
+    });
+
+    const runCuratorSpy = vi
+      .spyOn(ace as any, 'runCurator')
+      .mockResolvedValue(undefined);
+
+    const program = ax(
+      f().input('question', f.string()).output('answer', f.string()).build()
+    );
+    const examples = [
+      { question: 'q1', answer: 'a1' },
+      { question: 'q2', answer: 'a2' },
+    ];
+    const metricFn = () => 1;
+
+    await ace.compile(program, examples, metricFn, {
+      aceOptions: { maxEpochs: 1, maxReflectorRounds: 1 },
+    });
+
+    expect(runCuratorSpy).toHaveBeenCalled();
+    const curatorCall = runCuratorSpy.mock.calls[0][0];
+    expect(curatorCall.example).toHaveProperty('question');
+    expect(curatorCall.example).toHaveProperty('answer');
+  });
+});
+
 describe('Optimizer Interface', () => {
   it('AxBootstrapFewShot implements AxOptimizer interface', () => {
     const optimizer = new AxBootstrapFewShot({

src/ax/dsp/optimizers/ace.test.ts

@@ -1,4 +1,8 @@
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
+
+import { ax } from '../template.js';
+import { f } from '../sig.js';
+import type { AxAIService } from '../../ai/types';
 
 import { applyCuratorOperations, createEmptyPlaybook } from './acePlaybook.js';
 import { AxACE } from './ace.js';
@@ -128,3 +132,56 @@ describe('AxACE helpers', () => {
     expect(newBullet?.content).toBe('Third tactic');
   });
 });
+
+describe('AxACE', () => {
+  it('runCurator should only receive input fields in question_context', async () => {
+    const mockCuratorAI = {
+      name: 'mockCurator',
+      chat: vi.fn().mockResolvedValue({
+        results: [
+          {
+            index: 0,
+            content: '{"reasoning": "mock", "operations":[]}',
+          },
+        ],
+      }),
+      getOptions: () => ({ tracer: undefined }),
+      getLogger: () => undefined,
+    } as unknown as AxAIService;
+
+    const program = ax(
+      f().input('question', f.string()).output('answer', f.string()).build()
+    );
+
+    const example = {
+      question: 'This is the input',
+      answer: 'This is the output',
+    };
+
+    const ace = new AxACE({
+      studentAI: {} as any,
+      teacherAI: mockCuratorAI,
+    });
+
+    const curatorProgram = (ace as any).getOrCreateCuratorProgram();
+    const forwardSpy = vi.spyOn(curatorProgram, 'forward');
+
+    // Directly call the internal runCurator method for a focused unit test
+    await (ace as any).runCurator({
+      program,
+      example,
+      reflection: { keyInsight: 'test' }, // Minimal reflection to trigger curator
+      playbook: { sections: {}, stats: { bulletCount: 0 } },
+    });
+
+    expect(forwardSpy).toHaveBeenCalled();
+
+    const forwardArgs = forwardSpy.mock.calls[0][1] as any;
+    const receivedContext = JSON.parse(forwardArgs.question_context);
+
+    expect(receivedContext).toBeDefined();
+    expect(receivedContext).toHaveProperty('question');
+    expect(receivedContext.question).toBe('This is the input');
+    expect(receivedContext).not.toHaveProperty('answer');
+  });
+});

src/ax/dsp/optimizers/ace.ts

@@ -148,6 +148,8 @@ export class AxACE extends AxBaseOptimizer {
 
   private curatorProgram?: AxGen<any, any>;
 
+  private program?: Readonly<AxGen<any, any>>;
+
   constructor(
     args: Readonly<AxOptimizerArgs>,
     options?: Readonly<AxACEOptions>
@@ -220,6 +222,7 @@ export class AxACE extends AxBaseOptimizer {
 
     const startTime = Date.now();
     this.validateExamples(examples);
+    this.program = program;
 
     const baseInstruction = await this.extractProgramInstruction(program);
     const originalDescription = program.getSignature().getDescription() ?? '';
@@ -287,6 +290,7 @@ export class AxACE extends AxBaseOptimizer {
           });
 
           const rawCurator = await this.runCurator({
+            program,
             example,
             reflection,
             playbook: this.playbook,
@@ -458,6 +462,12 @@ export class AxACE extends AxBaseOptimizer {
       feedback?: string;
     }>
   ): Promise<AxACECuratorOutput | undefined> {
+    if (!this.program) {
+      throw new Error(
+        'AxACE: `compile` must be run before `applyOnlineUpdate`'
+      );
+    }
+
     const generatorOutput = this.createGeneratorOutput(
       args.prediction,
       args.example
@@ -480,6 +490,7 @@ export class AxACE extends AxBaseOptimizer {
     });
 
     const rawCurator = await this.runCurator({
+      program: this.program,
       example: args.example,
       reflection,
       playbook: this.playbook,
@@ -992,11 +1003,13 @@ export class AxACE extends AxBaseOptimizer {
     }
   }
 
-  private async runCurator({
+  private async runCurator<IN, OUT extends AxGenOut>({
+    program,
     example,
     reflection,
     playbook,
   }: Readonly<{
+    program: Readonly<AxGen<IN, OUT>>;
     example: AxExample;
     reflection?: AxACEReflectionOutput;
     playbook: AxACEPlaybook;
@@ -1008,14 +1021,26 @@ export class AxACE extends AxBaseOptimizer {
     const curator = this.getOrCreateCuratorProgram();
     const curatorAI = this.teacherAI ?? this.studentAI;
 
+    const signature = program.getSignature();
+    const inputFields = signature.getInputFields();
+    const questionContext = inputFields.reduce(
+      (acc, field) => {
+        if (field.name in example) {
+          acc[field.name] = example[field.name as keyof typeof example];
+        }
+        return acc;
+      },
+      {} as Record<string, unknown>
+    );
+
     try {
       const outputRaw = await curator.forward(curatorAI, {
         playbook: JSON.stringify({
           markdown: renderPlaybook(playbook),
           structured: playbook,
         }),
         reflection: JSON.stringify(reflection),
-        question_context: JSON.stringify(example),
+        question_context: JSON.stringify(questionContext),
         token_budget: 1024,
       });
 

src/ax/dsp/sigtypes-runtime.test.ts

@@ -69,10 +69,6 @@ describe('TypeScript Parser Parity with JS Parser', () => {
         expect(parsed.outputs.length).toBeGreaterThan(0);
       }).not.toThrow(`JS Parser failed for: "${signature}"`);
     }
-
-    console.log(
-      `✅ JS Parser successfully handled ${testCases.length} test cases`
-    );
   });
 
   test('should demonstrate TypeScript type inference capabilities', () => {
@@ -97,7 +93,5 @@ describe('TypeScript Parser Parity with JS Parser', () => {
     };
 
     expect(simpleClassExample.outputs.category).toBe('positive');
-
-    console.log('✅ TypeScript type inference working for supported cases');
   });
 });

src/ax/flow/flow-extensions.test.ts

@@ -1,12 +1,11 @@
 import { describe, expect, it } from 'vitest';
 
-import { AxFlow } from './flow.js';
+import { flow } from './flow.js';
 import { f } from '../dsp/sig.js';
 
 describe('AxFlow nodeExtended method', () => {
   it('should create chain-of-thought node with internal reasoning field', () => {
-    const flow = new AxFlow();
-    const cotFlow = flow.nodeExtended(
+    const cotFlow = flow().nodeExtended(
       'reasoner',
       'userInput:string -> answer:string',
       {
@@ -32,8 +31,7 @@ describe('AxFlow nodeExtended method', () => {
   });
 
   it('should create confidence-scoring node', () => {
-    const flow = new AxFlow();
-    const confFlow = flow.nodeExtended(
+    const confFlow = flow().nodeExtended(
       'scorer',
       'userInput:string -> analysis:string',
       {
@@ -56,8 +54,7 @@ describe('AxFlow nodeExtended method', () => {
   });
 
   it('should create contextual node with additional input fields', () => {
-    const flow = new AxFlow();
-    const contextFlow = flow.nodeExtended(
+    const contextFlow = flow().nodeExtended(
       'contextual',
       'question:string -> answer:string',
       {
@@ -82,8 +79,7 @@ describe('AxFlow nodeExtended method', () => {
   });
 
   it('should create extended node with all extension types', () => {
-    const flow = new AxFlow();
-    const extendedFlow = flow.nodeExtended(
+    const extendedFlow = flow().nodeExtended(
       'analyzer',
       'userInput:string -> analysis:string',
       {
@@ -130,10 +126,8 @@ describe('AxFlow nodeExtended method', () => {
   });
 
   it('should maintain type safety and prevent duplicate field names', () => {
-    const flow = new AxFlow();
-
     expect(() =>
-      flow.nodeExtended('test', 'userInput:string -> analysis:string', {
+      flow().nodeExtended('test', 'userInput:string -> analysis:string', {
         appendInputs: [
           {
             name: 'userInput',
@@ -145,10 +139,10 @@ describe('AxFlow nodeExtended method', () => {
   });
 
   it('should work with AxSignature instances as base', () => {
-    const flow = new AxFlow();
-    const baseSig = flow.getSignature(); // Get default signature
+    const myFlow = flow();
+    const baseSig = myFlow.getSignature(); // Get default signature
 
-    const extendedFlow = flow.nodeExtended('thinker', baseSig, {
+    const extendedFlow = myFlow.nodeExtended('thinker', baseSig, {
       prependOutputs: [
         {
           name: 'reasoning',
@@ -165,9 +159,7 @@ describe('AxFlow nodeExtended method', () => {
   });
 
   it('should support method chaining', () => {
-    const flow = new AxFlow();
-
-    const chainedFlow = flow
+    const chainedFlow = flow()
       .nodeExtended('reasoner', 'question:string -> analysis:string', {
         prependOutputs: [
           {
@@ -188,11 +180,9 @@ describe('AxFlow nodeExtended method', () => {
   });
 
   it('should validate field types according to input/output rules', () => {
-    const flow = new AxFlow();
-
     // Class types not allowed in input
     expect(() =>
-      flow.nodeExtended('test', 'userInput:string -> analysis:string', {
+      flow().nodeExtended('test', 'userInput:string -> analysis:string', {
         appendInputs: [
           {
             name: 'category',
@@ -204,7 +194,7 @@ describe('AxFlow nodeExtended method', () => {
 
     // Image types not allowed in output
     expect(() =>
-      flow.nodeExtended('test', 'userInput:string -> analysis:string', {
+      flow().nodeExtended('test', 'userInput:string -> analysis:string', {
         appendOutputs: [
           {
             name: 'outputImage',
@@ -216,10 +206,8 @@ describe('AxFlow nodeExtended method', () => {
   });
 
   it('should have nx alias that works identically to nodeExtended', () => {
-    const flow = new AxFlow();
-
     // Test nx alias with same functionality as nodeExtended
-    const nxFlow = flow.nx('reasoner', 'userInput:string -> answer:string', {
+    const nxFlow = flow().nx('reasoner', 'userInput:string -> answer:string', {
       prependOutputs: [
         {
           name: 'reasoning',