224 lines
8.8 KiB
TypeScript
224 lines
8.8 KiB
TypeScript
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);
|
|
});
|
|
});
|
|
});
|