moving

packages/optimizer/tests/optimizer.test.ts

@@ -0,0 +1,223 @@
+import { describe, it, expect, vi } from 'vitest';
+import {
+  SUMMARIZER_SIGNATURE,
+  CLASSIFIER_SIGNATURE,
+  PROMPT_EXPANDER_SIGNATURE,
+  createSummarizer,
+  createClassifier,
+  createPromptExpander,
+  getProgram,
+  PROGRAM_REGISTRY,
+  type ProgramName,
+} from '../src/signatures.js';
+import { PromptOptimizer } from '../src/optimizer.js';
+import type { AxAIService } from '@ax-llm/ax';
+
+// Convert camelCase to Title Case (what Ax uses in prompts)
+function toTitleCase(camel: string): string {
+  return camel.replace(/([A-Z])/g, ' $1').replace(/^./, s => s.toUpperCase()).trim();
+}
+
+// Mock AI service that returns predictable outputs
+// Must implement enough of AxAIService for AxGen.forward() to work
+function createMockAI(responses: Record<string, string>): AxAIService {
+  // Format response as Ax expects: Title Case field names followed by values
+  const content = Object.entries(responses)
+    .map(([k, v]) => `${toTitleCase(k)}: ${v}`)
+    .join('\n');
+
+  return {
+    chat: vi.fn().mockResolvedValue({
+      results: [{ content, index: 0 }],
+      modelUsage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 },
+    }),
+    embed: vi.fn(),
+    getFeatures: vi.fn().mockReturnValue({
+      functions: true,
+      streaming: true,
+      hasThinkingBudget: false,
+      hasShowThoughts: false,
+      structuredOutputs: false,
+      media: {},
+      caching: { supported: false, types: [] },
+      thinking: false,
+      multiTurn: true,
+    }),
+    getOptions: vi.fn().mockReturnValue({ debug: false, verbose: false }),
+    setOptions: vi.fn(),
+    getName: vi.fn().mockReturnValue('mock'),
+    getId: vi.fn().mockReturnValue('mock-id'),
+    getModelList: vi.fn().mockReturnValue([]),
+    getLastUsedChatModel: vi.fn().mockReturnValue('mock-model'),
+    getLastUsedEmbedModel: vi.fn().mockReturnValue(undefined),
+    getLastUsedModelConfig: vi.fn().mockReturnValue(undefined),
+    getMetrics: vi.fn().mockReturnValue(undefined),
+    getLogger: vi.fn().mockReturnValue(undefined),
+  } as unknown as AxAIService;
+}
+
+describe('Prompt Optimization (Ax Integration)', () => {
+  describe('Signature definitions', () => {
+    it('summarizer signature has correct input/output fields', () => {
+      const inputs = SUMMARIZER_SIGNATURE.getInputFields();
+      const outputs = SUMMARIZER_SIGNATURE.getOutputFields();
+      expect(inputs).toHaveLength(1);
+      expect(inputs[0].name).toBe('textToSummarize');
+      expect(outputs).toHaveLength(1);
+      expect(outputs[0].name).toBe('summaryText');
+    });
+
+    it('classifier signature has correct input/output fields', () => {
+      const inputs = CLASSIFIER_SIGNATURE.getInputFields();
+      const outputs = CLASSIFIER_SIGNATURE.getOutputFields();
+      expect(inputs).toHaveLength(1);
+      expect(inputs[0].name).toBe('textToClassify');
+      expect(outputs).toHaveLength(1);
+      expect(outputs[0].name).toBe('intentCategory');
+      expect(outputs[0].type?.name).toBe('class');
+      expect(outputs[0].type?.options).toEqual(['question', 'command', 'observation', 'request', 'greeting']);
+    });
+
+    it('prompt expander signature has correct input/output fields', () => {
+      const inputs = PROMPT_EXPANDER_SIGNATURE.getInputFields();
+      const outputs = PROMPT_EXPANDER_SIGNATURE.getOutputFields();
+      expect(inputs).toHaveLength(1);
+      expect(inputs[0].name).toBe('briefPrompt');
+      expect(outputs).toHaveLength(1);
+      expect(outputs[0].name).toBe('expandedPrompt');
+    });
+  });
+
+  describe('Program creation', () => {
+    it('creates summarizer program with instruction', () => {
+      const program = createSummarizer();
+      expect(program).toBeDefined();
+      expect(program.getInstruction()).toContain('summarizer');
+    });
+
+    it('creates classifier program with instruction', () => {
+      const program = createClassifier();
+      expect(program).toBeDefined();
+      expect(program.getInstruction()).toContain('Classify');
+    });
+
+    it('creates prompt expander program with instruction', () => {
+      const program = createPromptExpander();
+      expect(program).toBeDefined();
+      expect(program.getInstruction()).toContain('expand');
+    });
+  });
+
+  describe('Program registry', () => {
+    it('contains all three programs', () => {
+      expect(PROGRAM_REGISTRY).toHaveLength(3);
+      const names = PROGRAM_REGISTRY.map(p => p.name);
+      expect(names).toContain('summarizer');
+      expect(names).toContain('classifier');
+      expect(names).toContain('prompt_expander');
+    });
+
+    it('getProgram returns correct entry', () => {
+      const entry = getProgram('summarizer');
+      expect(entry.name).toBe('summarizer');
+      expect(typeof entry.create).toBe('function');
+      expect(entry.signature).toBe(SUMMARIZER_SIGNATURE);
+    });
+
+    it('getProgram throws for unknown program', () => {
+      expect(() => getProgram('nonexistent' as ProgramName)).toThrow('Unknown program: nonexistent');
+    });
+
+    it('each registry entry creates a valid program', () => {
+      for (const entry of PROGRAM_REGISTRY) {
+        const program = entry.create();
+        expect(program).toBeDefined();
+        expect(program.getInstruction()).toBeTruthy();
+      }
+    });
+  });
+
+  describe('PromptOptimizer', () => {
+    it('lists all available programs', () => {
+      const mockAI = createMockAI({});
+      const optimizer = new PromptOptimizer({ ai: mockAI });
+      const programs = optimizer.listPrograms();
+      expect(programs).toEqual(['summarizer', 'classifier', 'prompt_expander']);
+    });
+  });
+
+  // API-dependent tests - these test actual execution through the Ax framework
+  // They use a mock AI service, but still exercise the full AxGen.forward() pipeline
+  describe('Signature execution (API-dependent)', () => {
+    it('summarizer produces string output type', async () => {
+      const mockAI = createMockAI({ summaryText: 'This is a test summary.' });
+      const optimizer = new PromptOptimizer({ ai: mockAI });
+      const result = await optimizer.execute('summarizer', {
+        textToSummarize: 'A long text about AI agents and their capabilities in modern software.',
+      });
+      expect(result.programName).toBe('summarizer');
+      expect(result.input.textToSummarize).toContain('AI agents');
+      expect(typeof result.output.summaryText).toBe('string');
+    });
+
+    it('classifier returns valid category', async () => {
+      const mockAI = createMockAI({ intentCategory: 'question' });
+      const optimizer = new PromptOptimizer({ ai: mockAI });
+      const result = await optimizer.execute('classifier', {
+        textToClassify: 'What is the weather like today?',
+      });
+      expect(result.programName).toBe('classifier');
+      expect(['question', 'command', 'observation', 'request', 'greeting']).toContain(
+        result.output.intentCategory
+      );
+    });
+
+    it('prompt expander produces expanded output', async () => {
+      const mockAI = createMockAI({ expandedPrompt: 'Please write a comprehensive blog post about AI, covering...' });
+      const optimizer = new PromptOptimizer({ ai: mockAI });
+      const result = await optimizer.execute('prompt_expander', {
+        briefPrompt: 'Write about AI',
+      });
+      expect(result.programName).toBe('prompt_expander');
+      expect(typeof result.output.expandedPrompt).toBe('string');
+      expect(result.output.expandedPrompt.length).toBeGreaterThan(0);
+    });
+  });
+
+  describe('Intent classification across categories', () => {
+    const categories = [
+      { input: 'What time is it?', expected: 'question' },
+      { input: 'Delete all files now', expected: 'command' },
+      { input: 'The sky is blue', expected: 'observation' },
+      { input: 'Can you help me with this?', expected: 'request' },
+      { input: 'Hello, how are you?', expected: 'greeting' },
+    ];
+
+    for (const { input, expected } of categories) {
+      it(`classifies "${input}" as ${expected}`, async () => {
+        const mockAI = createMockAI({ intentCategory: expected });
+        const optimizer = new PromptOptimizer({ ai: mockAI });
+        const category = await optimizer.classify(input);
+        expect(category).toBe(expected);
+      });
+    }
+  });
+
+  describe('Convenience methods', () => {
+    it('summarize() returns summary string', async () => {
+      const mockAI = createMockAI({ summaryText: 'Brief summary here.' });
+      const optimizer = new PromptOptimizer({ ai: mockAI });
+      const summary = await optimizer.summarize('Long text about many topics...');
+      expect(typeof summary).toBe('string');
+      expect(summary.length).toBeGreaterThan(0);
+    });
+
+    it('expandPrompt() returns expanded string', async () => {
+      const mockAI = createMockAI({ expandedPrompt: 'Detailed expanded prompt...' });
+      const optimizer = new PromptOptimizer({ ai: mockAI });
+      const expanded = await optimizer.expandPrompt('Fix bug');
+      expect(typeof expanded).toBe('string');
+      expect(expanded.length).toBeGreaterThan(0);
+    });
+  });
+});