From c03b6813948fd338dd0580baf084b08a2d974f7e Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 4 Dec 2025 19:37:43 +0000 Subject: [PATCH] fix(ace): Refine reflector to use only input fields The AxACE optimizer's reflector was previously using the entire training example as the "question" for reflection, including output fields. This caused confusion, as the reflector was seeing both the inputs and the expected outputs, leading to inaccurate and unhelpful reflections. This commit fixes the issue by ensuring that only the input fields from an example are used as the "question" for the reflector. The code now retrieves the program's signature, identifies the defined input fields, and constructs a `questionContext` object containing only the input key-value pairs from the full example. By providing a clean separation between the input and output, the reflector can now generate more accurate and relevant insights, improving the overall effectiveness of the ACE optimization loop. --- src/ax/dsp/optimizers/ace.ts | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/ax/dsp/optimizers/ace.ts b/src/ax/dsp/optimizers/ace.ts index 87c6f756b..d0d2d3941 100644 --- a/src/ax/dsp/optimizers/ace.ts +++ b/src/ax/dsp/optimizers/ace.ts @@ -974,8 +974,20 @@ export class AxACE extends AxBaseOptimizer { severity: (example as { severity?: string })?.severity, policyHint: (example as { policyHint?: string })?.policyHint, }; + const signature = this.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 + ); + const reflectionRaw = await reflector.forward(reflectorAI, { - question: JSON.stringify(example), + question: JSON.stringify(questionContext), generator_answer: JSON.stringify(generatorOutput.answer), generator_reasoning: generatorOutput.reasoning, playbook: JSON.stringify({