diff --git a/src/ax/dsp/generate.ts b/src/ax/dsp/generate.ts index ba9b5f55d..6569b5b9b 100644 --- a/src/ax/dsp/generate.ts +++ b/src/ax/dsp/generate.ts @@ -180,6 +180,10 @@ export class AxGen this.promptTemplate.setInstruction(instruction); } + public getInstruction(): string | undefined { + return this.promptTemplate.getInstruction(); + } + private getSignatureName(): string { return this.signature.getDescription() || 'unknown_signature'; } diff --git a/src/ax/dsp/optimizers/gepa.test.ts b/src/ax/dsp/optimizers/gepa.test.ts new file mode 100644 index 000000000..9fd852a75 --- /dev/null +++ b/src/ax/dsp/optimizers/gepa.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it } from 'vitest'; +import { ax } from '../template.js'; +import type { AxAIService } from '../../ai/types.js'; +import { AxGEPA } from './gepa.js'; + +describe('AxGEPA Optimizer', () => { + describe('getBaseInstruction', () => { + it('should use the description from the signature when available', async () => { + // Create a program with a signature that has a description + const program = ax( + '"This is my custom task description" question:string -> answer:string' + ); + + // Access the private method via cast + const optimizer = new AxGEPA({ + studentAI: {} as AxAIService, + teacherAI: {} as AxAIService, + }); + + // Call getBaseInstruction + const instruction = await (optimizer as any).getBaseInstruction(program); + + // It should return the description from the signature, not the default + expect(instruction).toBe('This is my custom task description'); + expect(instruction).not.toBe( + 'Follow the task precisely. Be concise, correct, and consistent.' + ); + }); + + it('should fall back to default when signature has no description', async () => { + // Create a program without a description + const program = ax('question:string -> answer:string'); + + const optimizer = new AxGEPA({ + studentAI: {} as AxAIService, + teacherAI: {} as AxAIService, + }); + + const instruction = await (optimizer as any).getBaseInstruction(program); + + // Should use the default fallback + expect(instruction).toBe( + 'Follow the task precisely. Be concise, correct, and consistent.' + ); + }); + + it('should use custom instruction when set via setInstruction', async () => { + const program = ax('question:string -> answer:string'); + program.setInstruction('My explicitly set custom instruction'); + + const optimizer = new AxGEPA({ + studentAI: {} as AxAIService, + teacherAI: {} as AxAIService, + }); + + const instruction = await (optimizer as any).getBaseInstruction(program); + + // Should return the custom instruction + expect(instruction).toBe('My explicitly set custom instruction'); + }); + }); +}); diff --git a/src/ax/dsp/optimizers/gepa.ts b/src/ax/dsp/optimizers/gepa.ts index 59efcd964..206a57c9d 100644 --- a/src/ax/dsp/optimizers/gepa.ts +++ b/src/ax/dsp/optimizers/gepa.ts @@ -836,17 +836,19 @@ export class AxGEPA extends AxBaseOptimizer { private async getBaseInstruction( program: Readonly> ): Promise { - try { - // If program exposes instruction via signature, prefer it - const sig: any = program.getSignature?.(); - if ( - sig && - typeof sig.instruction === 'string' && - sig.instruction.length > 0 - ) { - return sig.instruction as string; - } - } catch {} + // First check for custom instruction set via setInstruction() + const customInstruction = program.getInstruction?.(); + if (customInstruction && customInstruction.length > 0) { + return customInstruction; + } + + // Fall back to signature description + const sig = program.getSignature?.(); + const description = sig?.getDescription?.(); + if (description && description.length > 0) { + return description; + } + return 'Follow the task precisely. Be concise, correct, and consistent.'; } diff --git a/src/ax/dsp/optimizers/gepaFlow.ts b/src/ax/dsp/optimizers/gepaFlow.ts index aff860dfd..50aa5f059 100644 --- a/src/ax/dsp/optimizers/gepaFlow.ts +++ b/src/ax/dsp/optimizers/gepaFlow.ts @@ -830,15 +830,19 @@ export class AxGEPAFlow extends AxBaseOptimizer { // === Helpers === private async getBaseInstruction(program: any): Promise { - try { - const sig = program?.getSignature?.(); - if ( - sig && - typeof sig.instruction === 'string' && - sig.instruction.length > 0 - ) - return sig.instruction as string; - } catch {} + // First check for custom instruction set via setInstruction() + const customInstruction = program?.getInstruction?.(); + if (customInstruction && customInstruction.length > 0) { + return customInstruction; + } + + // Fall back to signature description + const sig = program?.getSignature?.(); + const description = sig?.getDescription?.(); + if (description && description.length > 0) { + return description; + } + return 'Follow the task precisely. Be concise, correct, and consistent.'; } diff --git a/src/ax/dsp/prompt.ts b/src/ax/dsp/prompt.ts index 9b2072fa3..bbf554033 100644 --- a/src/ax/dsp/prompt.ts +++ b/src/ax/dsp/prompt.ts @@ -46,10 +46,16 @@ export class AxPromptTemplate { private sig: Readonly; private fieldTemplates?: Record; private task: { type: 'text'; text: string }; + private customInstruction?: string; public setInstruction(instruction: string): void { + this.customInstruction = instruction; this.task = { type: 'text', text: instruction }; } + + public getInstruction(): string | undefined { + return this.customInstruction; + } private readonly thoughtFieldName: string; private readonly functions?: Readonly; private readonly cacheSystemPrompt?: boolean;