532 lines
20 KiB
TypeScript
532 lines
20 KiB
TypeScript
import { describe, it, expect } from 'vitest';
|
||
import type { MemoryFrame, Importance, FrameType, FrameSource } from '@waggle/core';
|
||
import type { AgentPersona } from '../src/personas.js';
|
||
import type { TaskShape } from '../src/task-shape.js';
|
||
import type { ContextFrames } from '../src/orchestrator.js';
|
||
import {
|
||
PromptAssembler,
|
||
type AssembleInput,
|
||
type RecalledMemory,
|
||
} from '../src/prompt-assembler.js';
|
||
|
||
// ── Fixtures ─────────────────────────────────────────────────────────
|
||
|
||
let nextFrameId = 1;
|
||
function frame(
|
||
content: string,
|
||
opts: { type?: FrameType; importance?: Importance; source?: FrameSource } = {},
|
||
): MemoryFrame {
|
||
return {
|
||
id: nextFrameId++,
|
||
frame_type: opts.type ?? 'I',
|
||
gop_id: 'gop-1',
|
||
t: 0,
|
||
base_frame_id: null,
|
||
content,
|
||
importance: opts.importance ?? 'normal',
|
||
source: opts.source ?? 'user_stated',
|
||
access_count: 0,
|
||
created_at: '2026-04-17T10:00:00Z',
|
||
last_accessed: '2026-04-17T10:00:00Z',
|
||
};
|
||
}
|
||
|
||
function persona(overrides: Partial<AgentPersona> = {}): AgentPersona {
|
||
return {
|
||
id: 'researcher',
|
||
name: 'Researcher',
|
||
description: 'You are a researcher focused on evidence-based analysis.',
|
||
icon: 'search',
|
||
systemPrompt: '',
|
||
modelPreference: 'claude-opus-4-7',
|
||
tools: ['search_memory', 'save_memory'],
|
||
workspaceAffinity: [],
|
||
suggestedCommands: [],
|
||
defaultWorkflow: null,
|
||
tagline: 'Evidence-first thinker.',
|
||
...overrides,
|
||
};
|
||
}
|
||
|
||
function emptyContext(): ContextFrames {
|
||
return {
|
||
stateFrames: [],
|
||
recentChanges: [],
|
||
activeWork: [],
|
||
keyEntities: [],
|
||
personalPreferences: [],
|
||
};
|
||
}
|
||
|
||
function emptyRecalled(): RecalledMemory {
|
||
return { workspace: [], personal: [], scanSafe: true };
|
||
}
|
||
|
||
function shape(type: TaskShape['type'], confidence = 0.8): TaskShape {
|
||
return {
|
||
type,
|
||
confidence,
|
||
signals: [],
|
||
complexity: 'moderate',
|
||
};
|
||
}
|
||
|
||
function baseInput(overrides: Partial<AssembleInput> = {}): AssembleInput {
|
||
return {
|
||
corePrompt: 'You are Waggle, a memory-first AI colleague.',
|
||
persona: persona(),
|
||
context: emptyContext(),
|
||
recalled: emptyRecalled(),
|
||
query: 'test query',
|
||
tier: 'mid',
|
||
...overrides,
|
||
};
|
||
}
|
||
|
||
// ── Tests ────────────────────────────────────────────────────────────
|
||
|
||
describe('PromptAssembler.assemble', () => {
|
||
const assembler = new PromptAssembler();
|
||
|
||
it('includes Identity and Persona in every output', () => {
|
||
const out = assembler.assemble(baseInput());
|
||
expect(out.system).toContain('# Identity');
|
||
expect(out.system).toContain('## Persona: Researcher');
|
||
expect(out.debug.sectionsIncluded).toContain('Identity');
|
||
expect(out.debug.sectionsIncluded).toContain('Persona');
|
||
});
|
||
|
||
it('small tier caps State frames at 3', () => {
|
||
const frames: MemoryFrame[] = [];
|
||
for (let i = 0; i < 10; i++) {
|
||
frames.push(frame(`State frame ${i}`, { type: 'I', importance: 'important' }));
|
||
}
|
||
const out = assembler.assemble(
|
||
baseInput({
|
||
tier: 'small',
|
||
context: { ...emptyContext(), stateFrames: frames },
|
||
}),
|
||
);
|
||
const stateBlock = out.system.split('# State\n')[1]?.split('\n\n')[0] ?? '';
|
||
const stateLines = stateBlock.split('\n').filter(l => l.startsWith('-'));
|
||
expect(stateLines.length).toBeLessThanOrEqual(3);
|
||
});
|
||
|
||
it('frontier tier emits no Response-format section', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'frontier', taskShape: shape('plan-execute', 0.9) }),
|
||
);
|
||
expect(out.system).not.toContain('# Response format');
|
||
expect(out.responseScaffold).toBeNull();
|
||
expect(out.debug.scaffoldApplied).toBe(false);
|
||
});
|
||
|
||
it('mid + plan-execute yields "State plan. Execute. Report." scaffold', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'mid', taskShape: shape('plan-execute', 0.8) }),
|
||
);
|
||
expect(out.responseScaffold).toBe('State plan. Execute. Report.');
|
||
expect(out.system).toContain('# Response format');
|
||
expect(out.debug.scaffoldApplied).toBe(true);
|
||
});
|
||
|
||
it('small + compare yields the long assumption/trade-offs/recommendation scaffold', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.7) }),
|
||
);
|
||
expect(out.responseScaffold).toBe(
|
||
'State the assumption. List the trade-offs. Give the recommendation.',
|
||
);
|
||
});
|
||
|
||
it('draft shape emits no scaffold at any tier', () => {
|
||
for (const tier of ['small', 'mid', 'frontier'] as const) {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier, taskShape: shape('draft', 0.9) }),
|
||
);
|
||
expect(out.responseScaffold).toBeNull();
|
||
expect(out.debug.scaffoldApplied).toBe(false);
|
||
}
|
||
});
|
||
|
||
it('mixed shape emits no scaffold at any tier', () => {
|
||
for (const tier of ['small', 'mid', 'frontier'] as const) {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier, taskShape: shape('mixed', 0.9) }),
|
||
);
|
||
expect(out.responseScaffold).toBeNull();
|
||
}
|
||
});
|
||
|
||
it('confidence below 0.3 threshold: no scaffold even if shape would scaffold', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('research', 0.1) }),
|
||
);
|
||
expect(out.responseScaffold).toBeNull();
|
||
expect(out.debug.scaffoldApplied).toBe(false);
|
||
});
|
||
|
||
it('custom confidenceThreshold lets high-confidence scaffolds through', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('research', 0.2) }),
|
||
{ confidenceThreshold: 0.15 },
|
||
);
|
||
expect(out.responseScaffold).toBe(
|
||
'Cite the frame. Quote the relevant fragment. Answer directly.',
|
||
);
|
||
});
|
||
|
||
it('empty context.stateFrames: omits State section (no empty header)', () => {
|
||
const out = assembler.assemble(baseInput());
|
||
expect(out.system).not.toContain('# State');
|
||
expect(out.debug.sectionsIncluded).not.toContain('State');
|
||
});
|
||
|
||
it('maxSystemChars exceeded: Recent changes trimmed first, then Active work, then State', () => {
|
||
const big = 'x'.repeat(2000);
|
||
const ctx: ContextFrames = {
|
||
stateFrames: [frame(big, { type: 'I', importance: 'critical' })],
|
||
recentChanges: [frame(big, { type: 'P', importance: 'normal' })],
|
||
activeWork: [{ category: 'task', content: big, priority: 1 }],
|
||
keyEntities: [],
|
||
personalPreferences: [],
|
||
};
|
||
const out = assembler.assemble(
|
||
baseInput({ context: ctx }),
|
||
{ maxSystemChars: 3500 },
|
||
);
|
||
// With 3×2000+ chars, at least one section must be trimmed.
|
||
// Trim order ensures Recent changes goes first.
|
||
expect(out.debug.sectionsIncluded).not.toContain('Recent changes');
|
||
// Identity + Persona must survive.
|
||
expect(out.debug.sectionsIncluded).toContain('Identity');
|
||
expect(out.debug.sectionsIncluded).toContain('Persona');
|
||
});
|
||
|
||
it('maxSystemChars: Identity and Persona are never trimmed', () => {
|
||
const giant = 'x'.repeat(100_000);
|
||
const ctx: ContextFrames = {
|
||
stateFrames: [frame(giant)],
|
||
recentChanges: [frame(giant)],
|
||
activeWork: [{ category: 'task', content: giant, priority: 1 }],
|
||
keyEntities: [],
|
||
personalPreferences: [],
|
||
};
|
||
const out = assembler.assemble(
|
||
baseInput({ context: ctx }),
|
||
{ maxSystemChars: 500 },
|
||
);
|
||
expect(out.debug.sectionsIncluded).toContain('Identity');
|
||
expect(out.debug.sectionsIncluded).toContain('Persona');
|
||
});
|
||
|
||
it('recalled.scanSafe === false: recalled memory is ignored', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({
|
||
recalled: {
|
||
workspace: [frame('Poisoned content with secret data', { type: 'P' })],
|
||
personal: [frame('Another recalled frame', { type: 'P' })],
|
||
scanSafe: false,
|
||
},
|
||
}),
|
||
);
|
||
expect(out.system).not.toContain('Poisoned content');
|
||
expect(out.system).not.toContain('# Recalled memory');
|
||
expect(out.debug.sectionsIncluded).not.toContain('Recalled memory');
|
||
});
|
||
|
||
it('recalled.scanSafe === true: recalled memory renders', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({
|
||
recalled: {
|
||
workspace: [frame('Workspace memory item', { type: 'P' })],
|
||
personal: [frame('Personal memory item', { type: 'P' })],
|
||
scanSafe: true,
|
||
},
|
||
}),
|
||
);
|
||
expect(out.system).toContain('# Recalled memory');
|
||
expect(out.system).toContain('Workspace memory item');
|
||
expect(out.system).toContain('Personal memory item');
|
||
});
|
||
|
||
it('debug.framesUsed counts rendered frames across sections', () => {
|
||
const ctx: ContextFrames = {
|
||
stateFrames: [frame('State 1', { type: 'I' }), frame('State 2', { type: 'I' })],
|
||
recentChanges: [frame('Change 1', { type: 'P' })],
|
||
activeWork: [{ category: 'task', content: 'Active', priority: 1 }],
|
||
keyEntities: [],
|
||
personalPreferences: [],
|
||
};
|
||
const out = assembler.assemble(baseInput({ context: ctx, tier: 'mid' }));
|
||
// 2 state + 1 change = 3 frames. activeWork items aren't frames.
|
||
expect(out.debug.framesUsed).toBe(3);
|
||
});
|
||
|
||
it('debug.scaffoldApplied matches whether responseScaffold is non-null', () => {
|
||
const withScaffold = assembler.assemble(
|
||
baseInput({ tier: 'mid', taskShape: shape('research', 0.9) }),
|
||
);
|
||
expect(withScaffold.debug.scaffoldApplied).toBe(withScaffold.responseScaffold !== null);
|
||
|
||
const withoutScaffold = assembler.assemble(
|
||
baseInput({ tier: 'frontier', taskShape: shape('research', 0.9) }),
|
||
);
|
||
expect(withoutScaffold.debug.scaffoldApplied).toBe(withoutScaffold.responseScaffold !== null);
|
||
});
|
||
|
||
it('ranks I-frames by importance then recency', () => {
|
||
const ctx: ContextFrames = {
|
||
stateFrames: [
|
||
frame('Old normal', { type: 'I', importance: 'normal' }),
|
||
frame('New critical', { type: 'I', importance: 'critical' }),
|
||
frame('Middle important', { type: 'I', importance: 'important' }),
|
||
],
|
||
recentChanges: [],
|
||
activeWork: [],
|
||
keyEntities: [],
|
||
personalPreferences: [],
|
||
};
|
||
const out = assembler.assemble(baseInput({ context: ctx, tier: 'small' }));
|
||
const stateSection = out.system.split('# State\n')[1]?.split('\n\n')[0] ?? '';
|
||
const firstLine = stateSection.split('\n')[0];
|
||
expect(firstLine).toContain('critical');
|
||
expect(firstLine).toContain('New critical');
|
||
});
|
||
|
||
it('dedupes frames by id', () => {
|
||
const f = frame('Duplicate', { type: 'I' });
|
||
const ctx: ContextFrames = {
|
||
stateFrames: [f, f, f],
|
||
recentChanges: [],
|
||
activeWork: [],
|
||
keyEntities: [],
|
||
personalPreferences: [],
|
||
};
|
||
const out = assembler.assemble(baseInput({ context: ctx, tier: 'mid' }));
|
||
expect(out.debug.framesUsed).toBe(1);
|
||
});
|
||
|
||
it('debug.totalChars matches system.length exactly', () => {
|
||
const out = assembler.assemble(baseInput());
|
||
expect(out.debug.totalChars).toBe(out.system.length);
|
||
});
|
||
|
||
it('userPrefix is empty string in v1', () => {
|
||
const out = assembler.assemble(baseInput());
|
||
expect(out.userPrefix).toBe('');
|
||
});
|
||
|
||
it('tierOverride takes precedence over input.tier', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'frontier', taskShape: shape('research', 0.9) }),
|
||
{ tierOverride: 'small' },
|
||
);
|
||
expect(out.debug.tier).toBe('small');
|
||
// Small + research + high confidence → scaffold applies
|
||
expect(out.responseScaffold).not.toBeNull();
|
||
});
|
||
});
|
||
|
||
// ── v5: scaffoldStyle (compression vs expansion) ─────────────────────
|
||
//
|
||
// See PromptAssembler v5 brief §7.3, §12.2.
|
||
//
|
||
// v5 adds an EXPANSION variant alongside v4's COMPRESSION scaffolds.
|
||
// The hypothesis: dense instruction-tuned models (Gemma family) benefit
|
||
// from expansion scaffolds that give them more structure to fill, rather
|
||
// than compression scaffolds that tell them to say less.
|
||
|
||
describe('PromptAssembler.assemble — v5 scaffoldStyle', () => {
|
||
const assembler = new PromptAssembler();
|
||
|
||
it('default style equals explicit compression (v4 parity)', () => {
|
||
const noStyle = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
);
|
||
const explicit = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
{ scaffoldStyle: 'compression' },
|
||
);
|
||
expect(noStyle.responseScaffold).toBe(explicit.responseScaffold);
|
||
expect(noStyle.system).toBe(explicit.system);
|
||
});
|
||
|
||
it('compression + small + compare matches v4 text exactly (snapshot)', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
{ scaffoldStyle: 'compression' },
|
||
);
|
||
// Byte-identical to v4 (now COMPRESSION_SCAFFOLDS[compare][small]).
|
||
expect(out.responseScaffold).toBe(
|
||
'State the assumption. List the trade-offs. Give the recommendation.',
|
||
);
|
||
});
|
||
|
||
it('compression + mid + plan-execute matches v4 text exactly (snapshot)', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'mid', taskShape: shape('plan-execute', 0.8) }),
|
||
{ scaffoldStyle: 'compression' },
|
||
);
|
||
expect(out.responseScaffold).toBe('State plan. Execute. Report.');
|
||
});
|
||
|
||
it('compression + small + research matches v4 text exactly (snapshot)', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('research', 0.8) }),
|
||
{ scaffoldStyle: 'compression' },
|
||
);
|
||
expect(out.responseScaffold).toBe(
|
||
'Cite the frame. Quote the relevant fragment. Answer directly.',
|
||
);
|
||
});
|
||
|
||
it('expansion + small + compare includes Factors and Confidence keywords', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(out.responseScaffold).not.toBeNull();
|
||
expect(out.responseScaffold).toContain('Factors');
|
||
expect(out.responseScaffold).toContain('confidence');
|
||
expect(out.responseScaffold).toContain('Trade-offs');
|
||
expect(out.responseScaffold).toContain('Recommendation');
|
||
expect(out.responseScaffold).toContain('Assumption');
|
||
});
|
||
|
||
it('expansion + small + decide produces same template as compare (shared analysis block)', () => {
|
||
const compareOut = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
const decideOut = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('decide', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
const reviewOut = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('review', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(decideOut.responseScaffold).toBe(compareOut.responseScaffold);
|
||
expect(reviewOut.responseScaffold).toBe(compareOut.responseScaffold);
|
||
});
|
||
|
||
it('expansion + small + plan-execute uses named-phases template', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('plan-execute', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(out.responseScaffold).not.toBeNull();
|
||
expect(out.responseScaffold).toContain('named phases');
|
||
expect(out.responseScaffold).toContain('Goal');
|
||
expect(out.responseScaffold).toContain('Steps');
|
||
expect(out.responseScaffold).toContain('Dependencies');
|
||
expect(out.responseScaffold).toContain('Blockers');
|
||
expect(out.responseScaffold).toContain('timeline');
|
||
expect(out.responseScaffold).toContain('critical-path');
|
||
});
|
||
|
||
it('expansion + small + research uses three-part template', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('research', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(out.responseScaffold).not.toBeNull();
|
||
expect(out.responseScaffold).toContain('Direct answer');
|
||
expect(out.responseScaffold).toContain('Source');
|
||
expect(out.responseScaffold).toContain('Context');
|
||
expect(out.responseScaffold).toContain('three parts');
|
||
});
|
||
|
||
it('expansion + mid produces shorter variant than small for same shape', () => {
|
||
const small = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
const mid = assembler.assemble(
|
||
baseInput({ tier: 'mid', taskShape: shape('compare', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(mid.responseScaffold).not.toBeNull();
|
||
expect(mid.responseScaffold!.length).toBeLessThan(small.responseScaffold!.length);
|
||
});
|
||
|
||
it('expansion + draft shape → null at every tier (creative never scaffolded)', () => {
|
||
for (const tier of ['small', 'mid', 'frontier'] as const) {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier, taskShape: shape('draft', 0.9) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(out.responseScaffold).toBeNull();
|
||
expect(out.debug.scaffoldApplied).toBe(false);
|
||
}
|
||
});
|
||
|
||
it('expansion + mixed shape → null at every tier', () => {
|
||
for (const tier of ['small', 'mid', 'frontier'] as const) {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier, taskShape: shape('mixed', 0.9) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(out.responseScaffold).toBeNull();
|
||
}
|
||
});
|
||
|
||
it('expansion + frontier → null (expansion forced to compression at frontier)', () => {
|
||
// Per §7.3: "Frontier tier receives compression scaffold regardless of
|
||
// scaffoldStyle request." Since compression[*][frontier] = null in v4,
|
||
// expansion at frontier also yields null. Don't risk v4's F > E
|
||
// replication by applying expansion to frontier models.
|
||
for (const shapeType of ['compare', 'decide', 'review', 'plan-execute', 'research'] as const) {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'frontier', taskShape: shape(shapeType, 0.9) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(out.responseScaffold).toBeNull();
|
||
}
|
||
});
|
||
|
||
it('expansion respects confidence threshold (same gate as compression)', () => {
|
||
const out = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.1) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(out.responseScaffold).toBeNull();
|
||
});
|
||
|
||
it('debug.scaffoldStyle reflects the requested style', () => {
|
||
const defaultOut = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
);
|
||
expect(defaultOut.debug.scaffoldStyle).toBe('compression');
|
||
|
||
const compressionOut = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
{ scaffoldStyle: 'compression' },
|
||
);
|
||
expect(compressionOut.debug.scaffoldStyle).toBe('compression');
|
||
|
||
const expansionOut = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(expansionOut.debug.scaffoldStyle).toBe('expansion');
|
||
});
|
||
|
||
it('debug.scaffoldStyle populated even when no scaffold is emitted (frontier / draft / low-conf)', () => {
|
||
const frontierOut = assembler.assemble(
|
||
baseInput({ tier: 'frontier', taskShape: shape('compare', 0.9) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(frontierOut.debug.scaffoldStyle).toBe('expansion');
|
||
expect(frontierOut.debug.scaffoldApplied).toBe(false);
|
||
|
||
const draftOut = assembler.assemble(
|
||
baseInput({ tier: 'small', taskShape: shape('draft', 0.9) }),
|
||
{ scaffoldStyle: 'expansion' },
|
||
);
|
||
expect(draftOut.debug.scaffoldStyle).toBe('expansion');
|
||
expect(draftOut.debug.scaffoldApplied).toBe(false);
|
||
});
|
||
});
|