329 lines
12 KiB
TypeScript
329 lines
12 KiB
TypeScript
/**
|
||
* Task 2.5 Stage 2-Retry §1.4 — agent-loop tool-exhaustion fallback tests.
|
||
*
|
||
* Stage 2 N=20 showed 2/20 agentic instances reached maxTurns=3 with every
|
||
* turn spent on a search_memory call, leaving `resp.content` empty; the
|
||
* judge scored those as incorrect. Stage 2-Retry §1.4 adds a runtime-side
|
||
* forced-answer fallback in the agentic cell wrapper: on empty-content +
|
||
* non-empty toolsUsed, the cell makes ONE additional direct LLM call
|
||
* (no tools, SYSTEM_AGENTIC_FORCED_FALLBACK) with the accumulated search
|
||
* context and returns that answer.
|
||
*
|
||
* Test matrix per brief:
|
||
* (a) normal 1-call-1-answer → fallback NOT fired
|
||
* (b) 2-call-1-answer → fallback NOT fired
|
||
* (c) 3-call-1-answer → fallback NOT fired
|
||
* (d) 3 tools + empty content → fallback FIRED, returns forced answer
|
||
* Plus:
|
||
* (e) fallback preserves accumulated tool context in the user prompt
|
||
* (f) fallback counts its tokens into returned cost (no cost-leak)
|
||
*/
|
||
|
||
import { describe, expect, it } from 'vitest';
|
||
import type { AgentLoopConfig, AgentResponse } from '@waggle/agent';
|
||
import type { LlmCallInput, LlmCallResult, LlmClient } from '../src/llm.js';
|
||
import type { DatasetInstance, ModelSpec } from '../src/types.js';
|
||
import {
|
||
cells,
|
||
SYSTEM_AGENTIC,
|
||
SYSTEM_AGENTIC_FORCED_FALLBACK,
|
||
} from '../src/cells.js';
|
||
import { createSubstrate } from '../src/substrate.js';
|
||
import type { Embedder } from '@waggle/core';
|
||
|
||
const VEC_DIMS = 1024;
|
||
function createFakeEmbedder(): Embedder {
|
||
// Deterministic hash-seeded 1024-dim embedder — same pattern as other tests.
|
||
const fnv1a = (s: string): number => {
|
||
let h = 2166136261 >>> 0;
|
||
for (let i = 0; i < s.length; i++) {
|
||
h ^= s.charCodeAt(i);
|
||
h = Math.imul(h, 16777619) >>> 0;
|
||
}
|
||
return h || 1;
|
||
};
|
||
const embedOne = (text: string): Float32Array => {
|
||
let state = fnv1a(text);
|
||
const v = new Float32Array(VEC_DIMS);
|
||
for (let i = 0; i < VEC_DIMS; i++) {
|
||
state ^= state << 13; state >>>= 0;
|
||
state ^= state >>> 17;
|
||
state ^= state << 5; state >>>= 0;
|
||
v[i] = ((state >>> 0) / 0x100000000) * 2 - 1;
|
||
}
|
||
let mag = 0;
|
||
for (let i = 0; i < VEC_DIMS; i++) mag += v[i] * v[i];
|
||
mag = Math.sqrt(mag);
|
||
if (mag > 0) for (let i = 0; i < VEC_DIMS; i++) v[i] /= mag;
|
||
return v;
|
||
};
|
||
return {
|
||
dimensions: VEC_DIMS,
|
||
async embed(t) { return embedOne(t); },
|
||
async embedBatch(ts) { return ts.map(embedOne); },
|
||
};
|
||
}
|
||
|
||
const MODEL: ModelSpec = {
|
||
id: 'test-subject',
|
||
displayName: 'Test',
|
||
provider: 'alibaba',
|
||
litellmModel: 'test/model',
|
||
pricePerMillionInput: 0.1,
|
||
pricePerMillionOutput: 0.4,
|
||
contextWindow: 32_000,
|
||
};
|
||
const INSTANCE: DatasetInstance = {
|
||
instance_id: 'test_q001',
|
||
question: 'When did the event happen?',
|
||
context: 'irrelevant for this test',
|
||
expected: ['2023'],
|
||
conversation_id: 'conv-test',
|
||
};
|
||
|
||
/** Programmable mock runAgentLoop that simulates different tool-use and
|
||
* content patterns per-test. Also exercises the onToolResult callback so
|
||
* the cell wrapper's context capture is tested end-to-end. */
|
||
function makeScriptedAgentLoop(script: {
|
||
toolCallResults: string[]; // results the agent-loop would return to the agent
|
||
finalContent: string; // final resp.content
|
||
usageInput?: number;
|
||
usageOutput?: number;
|
||
}): (config: AgentLoopConfig) => Promise<AgentResponse> {
|
||
return async (config: AgentLoopConfig): Promise<AgentResponse> => {
|
||
// Fire onToolResult for each simulated tool call in order. Triggers the
|
||
// cell wrapper's capturedToolResults accumulator.
|
||
for (let i = 0; i < script.toolCallResults.length; i++) {
|
||
config.onToolResult?.('search_memory', { query: `simulated-${i}` }, script.toolCallResults[i]);
|
||
}
|
||
return {
|
||
content: script.finalContent,
|
||
toolsUsed: script.toolCallResults.map(() => 'search_memory'),
|
||
usage: {
|
||
inputTokens: script.usageInput ?? 200,
|
||
outputTokens: script.usageOutput ?? 10,
|
||
},
|
||
};
|
||
};
|
||
}
|
||
|
||
/** Capturing LlmClient — records every direct llm.call made by the cell
|
||
* wrapper. The forced-fallback pass hits this (not the mock agent-loop). */
|
||
function makeCapturingLlm(response: Partial<LlmCallResult> = {}): {
|
||
client: LlmClient;
|
||
calls: LlmCallInput[];
|
||
} {
|
||
const calls: LlmCallInput[] = [];
|
||
const client: LlmClient = {
|
||
async call(input) {
|
||
calls.push(input);
|
||
return {
|
||
text: response.text ?? 'FORCED_ANSWER',
|
||
inputTokens: response.inputTokens ?? 500,
|
||
outputTokens: response.outputTokens ?? 3,
|
||
latencyMs: response.latencyMs ?? 50,
|
||
costUsd: response.costUsd ?? 0.0002,
|
||
failureMode: response.failureMode ?? null,
|
||
};
|
||
},
|
||
};
|
||
return { client, calls };
|
||
}
|
||
|
||
describe('agent-loop tool-exhaustion fallback — Stage 2-Retry §1.4', () => {
|
||
it('case (a): 1 tool call + answer → fallback NOT fired', async () => {
|
||
const substrate = createSubstrate({ embedder: createFakeEmbedder() });
|
||
try {
|
||
const runFn = makeScriptedAgentLoop({
|
||
toolCallResults: ['[1] Caroline: painted in 2023'],
|
||
finalContent: '2023',
|
||
});
|
||
const { client, calls } = makeCapturingLlm();
|
||
const result = await cells.agentic({
|
||
instance: INSTANCE, model: MODEL, llm: client, turnId: 't',
|
||
substrate, litellm: { url: 'u', apiKey: 'k' }, runAgentLoopFn: runFn,
|
||
});
|
||
expect(result.text).toBe('2023');
|
||
expect(result.failureMode).toBeNull();
|
||
// No fallback call.
|
||
expect(calls).toHaveLength(0);
|
||
} finally {
|
||
substrate.close();
|
||
}
|
||
});
|
||
|
||
it('case (b): 2 tool calls + answer → fallback NOT fired', async () => {
|
||
const substrate = createSubstrate({ embedder: createFakeEmbedder() });
|
||
try {
|
||
const runFn = makeScriptedAgentLoop({
|
||
toolCallResults: [
|
||
'[1] some result',
|
||
'[2] refined result',
|
||
],
|
||
finalContent: '2023',
|
||
});
|
||
const { client, calls } = makeCapturingLlm();
|
||
const result = await cells.agentic({
|
||
instance: INSTANCE, model: MODEL, llm: client, turnId: 't',
|
||
substrate, litellm: { url: 'u', apiKey: 'k' }, runAgentLoopFn: runFn,
|
||
});
|
||
expect(result.text).toBe('2023');
|
||
expect(calls).toHaveLength(0);
|
||
} finally {
|
||
substrate.close();
|
||
}
|
||
});
|
||
|
||
it('case (c): 3 tool calls + answer (all turns used, but content present) → fallback NOT fired', async () => {
|
||
const substrate = createSubstrate({ embedder: createFakeEmbedder() });
|
||
try {
|
||
const runFn = makeScriptedAgentLoop({
|
||
toolCallResults: ['r1', 'r2', 'r3'],
|
||
finalContent: 'best-effort-answer',
|
||
});
|
||
const { client, calls } = makeCapturingLlm();
|
||
const result = await cells.agentic({
|
||
instance: INSTANCE, model: MODEL, llm: client, turnId: 't',
|
||
substrate, litellm: { url: 'u', apiKey: 'k' }, runAgentLoopFn: runFn,
|
||
});
|
||
expect(result.text).toBe('best-effort-answer');
|
||
expect(calls).toHaveLength(0);
|
||
} finally {
|
||
substrate.close();
|
||
}
|
||
});
|
||
|
||
it('case (d): 3 tool calls + EMPTY content → fallback FIRED, forced answer returned', async () => {
|
||
const substrate = createSubstrate({ embedder: createFakeEmbedder() });
|
||
try {
|
||
const runFn = makeScriptedAgentLoop({
|
||
toolCallResults: ['r1', 'r2', 'r3'],
|
||
finalContent: '', // agent exhausted turns, no answer
|
||
});
|
||
const { client, calls } = makeCapturingLlm({ text: 'FALLBACK-2023' });
|
||
const result = await cells.agentic({
|
||
instance: INSTANCE, model: MODEL, llm: client, turnId: 't',
|
||
substrate, litellm: { url: 'u', apiKey: 'k' }, runAgentLoopFn: runFn,
|
||
});
|
||
expect(result.text).toBe('FALLBACK-2023');
|
||
expect(result.failureMode).toBeNull();
|
||
// Exactly one fallback call.
|
||
expect(calls).toHaveLength(1);
|
||
expect(calls[0].systemPrompt).toBe(SYSTEM_AGENTIC_FORCED_FALLBACK);
|
||
} finally {
|
||
substrate.close();
|
||
}
|
||
});
|
||
|
||
it('case (e): fallback user prompt includes question + every captured tool result', async () => {
|
||
const substrate = createSubstrate({ embedder: createFakeEmbedder() });
|
||
try {
|
||
const runFn = makeScriptedAgentLoop({
|
||
toolCallResults: [
|
||
'search_memory hit A',
|
||
'search_memory hit B',
|
||
'search_memory hit C',
|
||
],
|
||
finalContent: '',
|
||
});
|
||
const { client, calls } = makeCapturingLlm({ text: 'forced' });
|
||
await cells.agentic({
|
||
instance: INSTANCE, model: MODEL, llm: client, turnId: 't',
|
||
substrate, litellm: { url: 'u', apiKey: 'k' }, runAgentLoopFn: runFn,
|
||
});
|
||
const user = calls[0].userPrompt;
|
||
expect(user).toContain(INSTANCE.question);
|
||
expect(user).toContain('search_memory hit A');
|
||
expect(user).toContain('search_memory hit B');
|
||
expect(user).toContain('search_memory hit C');
|
||
// Each call result labelled with its call number.
|
||
expect(user).toContain('## search_memory call 1');
|
||
expect(user).toContain('## search_memory call 2');
|
||
expect(user).toContain('## search_memory call 3');
|
||
} finally {
|
||
substrate.close();
|
||
}
|
||
});
|
||
|
||
it('case (f): fallback token counts fold into the returned LlmCallResult cost', async () => {
|
||
const substrate = createSubstrate({ embedder: createFakeEmbedder() });
|
||
try {
|
||
const runFn = makeScriptedAgentLoop({
|
||
toolCallResults: ['r1', 'r2', 'r3'],
|
||
finalContent: '',
|
||
usageInput: 1_000_000,
|
||
usageOutput: 100_000,
|
||
});
|
||
// Fallback call accounts for another 500_000 input + 50_000 output.
|
||
const { client, calls } = makeCapturingLlm({
|
||
text: 'forced-final',
|
||
inputTokens: 500_000,
|
||
outputTokens: 50_000,
|
||
});
|
||
const result = await cells.agentic({
|
||
instance: INSTANCE, model: MODEL, llm: client, turnId: 't',
|
||
substrate, litellm: { url: 'u', apiKey: 'k' }, runAgentLoopFn: runFn,
|
||
});
|
||
// Token sums: 1_500_000 input, 150_000 output.
|
||
expect(result.inputTokens).toBe(1_500_000);
|
||
expect(result.outputTokens).toBe(150_000);
|
||
// Cost = (1.5 × $0.1/M input) + (0.15 × $0.4/M output) = $0.15 + $0.06 = $0.21
|
||
expect(result.costUsd).toBeCloseTo(0.21, 5);
|
||
// And one fallback llm.call happened.
|
||
expect(calls).toHaveLength(1);
|
||
} finally {
|
||
substrate.close();
|
||
}
|
||
});
|
||
|
||
it('case (g): empty content + ZERO tool calls → fallback does NOT fire (honest abstain)', async () => {
|
||
const substrate = createSubstrate({ embedder: createFakeEmbedder() });
|
||
try {
|
||
// Agent answered directly without searching — and its content was empty.
|
||
// This is an honest abstain case, not tool-exhaustion; no fallback.
|
||
const runFn = makeScriptedAgentLoop({
|
||
toolCallResults: [],
|
||
finalContent: '',
|
||
});
|
||
const { client, calls } = makeCapturingLlm();
|
||
const result = await cells.agentic({
|
||
instance: INSTANCE, model: MODEL, llm: client, turnId: 't',
|
||
substrate, litellm: { url: 'u', apiKey: 'k' }, runAgentLoopFn: runFn,
|
||
});
|
||
expect(result.text).toBe('');
|
||
expect(calls).toHaveLength(0);
|
||
} finally {
|
||
substrate.close();
|
||
}
|
||
});
|
||
|
||
it('SYSTEM_AGENTIC_FORCED_FALLBACK is exported and non-empty', () => {
|
||
expect(SYSTEM_AGENTIC_FORCED_FALLBACK.length).toBeGreaterThan(50);
|
||
expect(SYSTEM_AGENTIC_FORCED_FALLBACK).toContain('commit to your best');
|
||
expect(SYSTEM_AGENTIC_FORCED_FALLBACK).toContain('Do not call tools');
|
||
});
|
||
|
||
it('SYSTEM_AGENTIC (softened Stage 2-Retry) is still the prompt agentic uses', async () => {
|
||
const substrate = createSubstrate({ embedder: createFakeEmbedder() });
|
||
try {
|
||
let capturedSystem = '';
|
||
const runFn = async (config: AgentLoopConfig): Promise<AgentResponse> => {
|
||
capturedSystem = config.systemPrompt;
|
||
return { content: 'x', toolsUsed: [], usage: { inputTokens: 0, outputTokens: 0 } };
|
||
};
|
||
const { client } = makeCapturingLlm();
|
||
await cells.agentic({
|
||
instance: INSTANCE, model: MODEL, llm: client, turnId: 't',
|
||
substrate, litellm: { url: 'u', apiKey: 'k' }, runAgentLoopFn: runFn,
|
||
});
|
||
expect(capturedSystem).toBe(SYSTEM_AGENTIC);
|
||
// Sanity check on the softened surface language that diverges from Stage 1.
|
||
expect(capturedSystem).toContain('Protocol (you SHOULD follow)');
|
||
expect(capturedSystem).not.toContain('Protocol (you MUST follow)');
|
||
expect(capturedSystem).toContain('general knowledge');
|
||
} finally {
|
||
substrate.close();
|
||
}
|
||
});
|
||
});
|