This commit is contained in:
Oleg Maslov
280
packages/agent/tests/optimization-capture.test.ts
Normal file
280
packages/agent/tests/optimization-capture.test.ts
Normal file
@@ -0,0 +1,280 @@
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { MindDB, OptimizationLogStore } from '@waggle/core';
|
||||
import {
|
||||
captureInteraction,
|
||||
getRecentLogs,
|
||||
getWorkspaceLogs,
|
||||
isWithinBudget,
|
||||
} from '../src/optimization-capture.js';
|
||||
import type { AgentResponse } from '../src/agent-loop.js';
|
||||
|
||||
describe('optimization-capture', () => {
|
||||
let db: MindDB;
|
||||
let store: OptimizationLogStore;
|
||||
|
||||
beforeEach(() => {
|
||||
db = new MindDB(':memory:');
|
||||
store = new OptimizationLogStore(db);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
db.close();
|
||||
});
|
||||
|
||||
function makeResponse(overrides?: Partial<AgentResponse>): AgentResponse {
|
||||
return {
|
||||
content: 'Hello, I can help with that.',
|
||||
toolsUsed: ['search_memory', 'save_memory'],
|
||||
usage: { inputTokens: 500, outputTokens: 200 },
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
// ── captureInteraction ─────────────────────────────────────
|
||||
|
||||
describe('captureInteraction', () => {
|
||||
it('stores a log entry with correct fields', () => {
|
||||
const entry = captureInteraction(store, {
|
||||
sessionId: 'session-1',
|
||||
workspaceId: 'ws-abc',
|
||||
systemPrompt: 'You are a helpful assistant.',
|
||||
response: makeResponse(),
|
||||
turnCount: 3,
|
||||
wasCorrection: false,
|
||||
});
|
||||
|
||||
expect(entry.id).toBeGreaterThan(0);
|
||||
expect(entry.session_id).toBe('session-1');
|
||||
expect(entry.workspace_id).toBe('ws-abc');
|
||||
expect(entry.system_prompt).toBe('You are a helpful assistant.');
|
||||
expect(JSON.parse(entry.tools_used)).toEqual(['search_memory', 'save_memory']);
|
||||
expect(entry.turn_count).toBe(3);
|
||||
expect(entry.was_correction).toBe(0);
|
||||
expect(entry.input_tokens).toBe(500);
|
||||
expect(entry.output_tokens).toBe(200);
|
||||
expect(entry.timestamp).toBeTruthy();
|
||||
});
|
||||
|
||||
it('records wasCorrection as 1 when true', () => {
|
||||
const entry = captureInteraction(store, {
|
||||
sessionId: 'session-1',
|
||||
workspaceId: 'ws-abc',
|
||||
systemPrompt: 'test prompt',
|
||||
response: makeResponse(),
|
||||
turnCount: 1,
|
||||
wasCorrection: true,
|
||||
});
|
||||
|
||||
expect(entry.was_correction).toBe(1);
|
||||
});
|
||||
|
||||
it('handles empty toolsUsed', () => {
|
||||
const entry = captureInteraction(store, {
|
||||
sessionId: 'session-1',
|
||||
workspaceId: 'ws-abc',
|
||||
systemPrompt: 'test prompt',
|
||||
response: makeResponse({ toolsUsed: [] }),
|
||||
turnCount: 1,
|
||||
wasCorrection: false,
|
||||
});
|
||||
|
||||
expect(JSON.parse(entry.tools_used)).toEqual([]);
|
||||
});
|
||||
|
||||
it('defaults token counts to 0 when missing', () => {
|
||||
const response: AgentResponse = {
|
||||
content: 'response',
|
||||
toolsUsed: [],
|
||||
usage: { inputTokens: 0, outputTokens: 0 },
|
||||
};
|
||||
const entry = captureInteraction(store, {
|
||||
sessionId: 's1',
|
||||
workspaceId: 'ws1',
|
||||
systemPrompt: 'prompt',
|
||||
response,
|
||||
turnCount: 1,
|
||||
wasCorrection: false,
|
||||
});
|
||||
|
||||
expect(entry.input_tokens).toBe(0);
|
||||
expect(entry.output_tokens).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
// ── getRecentLogs ──────────────────────────────────────────
|
||||
|
||||
describe('getRecentLogs', () => {
|
||||
it('retrieves stored entries ordered by most recent first', () => {
|
||||
captureInteraction(store, {
|
||||
sessionId: 's1', workspaceId: 'ws1', systemPrompt: 'p1',
|
||||
response: makeResponse(), turnCount: 1, wasCorrection: false,
|
||||
});
|
||||
captureInteraction(store, {
|
||||
sessionId: 's2', workspaceId: 'ws1', systemPrompt: 'p2',
|
||||
response: makeResponse(), turnCount: 2, wasCorrection: true,
|
||||
});
|
||||
|
||||
const logs = getRecentLogs(store, 10);
|
||||
expect(logs).toHaveLength(2);
|
||||
// Ordered by timestamp DESC, id DESC — s2 was inserted second so it comes first
|
||||
expect(logs[0].session_id).toBe('s2');
|
||||
expect(logs[1].session_id).toBe('s1');
|
||||
});
|
||||
|
||||
it('respects the limit parameter', () => {
|
||||
for (let i = 0; i < 5; i++) {
|
||||
captureInteraction(store, {
|
||||
sessionId: `s${i}`, workspaceId: 'ws1', systemPrompt: 'p',
|
||||
response: makeResponse(), turnCount: 1, wasCorrection: false,
|
||||
});
|
||||
}
|
||||
|
||||
const logs = getRecentLogs(store, 3);
|
||||
expect(logs).toHaveLength(3);
|
||||
});
|
||||
|
||||
it('returns empty array when no entries exist', () => {
|
||||
const logs = getRecentLogs(store);
|
||||
expect(logs).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
// ── getWorkspaceLogs ───────────────────────────────────────
|
||||
|
||||
describe('getWorkspaceLogs', () => {
|
||||
it('filters by workspace ID', () => {
|
||||
captureInteraction(store, {
|
||||
sessionId: 's1', workspaceId: 'ws-a', systemPrompt: 'p',
|
||||
response: makeResponse(), turnCount: 1, wasCorrection: false,
|
||||
});
|
||||
captureInteraction(store, {
|
||||
sessionId: 's2', workspaceId: 'ws-b', systemPrompt: 'p',
|
||||
response: makeResponse(), turnCount: 1, wasCorrection: false,
|
||||
});
|
||||
captureInteraction(store, {
|
||||
sessionId: 's3', workspaceId: 'ws-a', systemPrompt: 'p',
|
||||
response: makeResponse(), turnCount: 2, wasCorrection: true,
|
||||
});
|
||||
|
||||
const wsALogs = getWorkspaceLogs(store, 'ws-a');
|
||||
expect(wsALogs).toHaveLength(2);
|
||||
expect(wsALogs.every(l => l.workspace_id === 'ws-a')).toBe(true);
|
||||
|
||||
const wsBLogs = getWorkspaceLogs(store, 'ws-b');
|
||||
expect(wsBLogs).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
|
||||
// ── Entries have correct shape ─────────────────────────────
|
||||
|
||||
describe('entry shape', () => {
|
||||
it('has all required fields with correct types', () => {
|
||||
const entry = captureInteraction(store, {
|
||||
sessionId: 'session-x',
|
||||
workspaceId: 'workspace-y',
|
||||
systemPrompt: 'You are helpful.',
|
||||
response: makeResponse({ toolsUsed: ['bash', 'read_file'] }),
|
||||
turnCount: 7,
|
||||
wasCorrection: false,
|
||||
});
|
||||
|
||||
// Check all fields exist and are of expected types
|
||||
expect(typeof entry.id).toBe('number');
|
||||
expect(typeof entry.session_id).toBe('string');
|
||||
expect(typeof entry.workspace_id).toBe('string');
|
||||
expect(typeof entry.system_prompt).toBe('string');
|
||||
expect(typeof entry.tools_used).toBe('string');
|
||||
expect(typeof entry.turn_count).toBe('number');
|
||||
expect(typeof entry.was_correction).toBe('number');
|
||||
expect(typeof entry.input_tokens).toBe('number');
|
||||
expect(typeof entry.output_tokens).toBe('number');
|
||||
expect(typeof entry.timestamp).toBe('string');
|
||||
|
||||
// tools_used is valid JSON array
|
||||
const parsed = JSON.parse(entry.tools_used);
|
||||
expect(Array.isArray(parsed)).toBe(true);
|
||||
expect(parsed).toEqual(['bash', 'read_file']);
|
||||
});
|
||||
});
|
||||
|
||||
// ── Budget enforcement ─────────────────────────────────────
|
||||
|
||||
describe('isWithinBudget', () => {
|
||||
it('returns true when no entries exist (under budget)', () => {
|
||||
expect(isWithinBudget(store, 100)).toBe(true);
|
||||
});
|
||||
|
||||
it('returns true when token cost is under budget', () => {
|
||||
// Insert entry with small token count
|
||||
captureInteraction(store, {
|
||||
sessionId: 's1', workspaceId: 'ws1', systemPrompt: 'p',
|
||||
response: makeResponse({ usage: { inputTokens: 100, outputTokens: 50 } }),
|
||||
turnCount: 1, wasCorrection: false,
|
||||
});
|
||||
|
||||
// 100 input * $3/M + 50 output * $15/M = $0.0003 + $0.00075 = ~$0.00105 = ~0.1 cents
|
||||
// Budget is 100 cents = $1, so well within budget
|
||||
expect(isWithinBudget(store, 100)).toBe(true);
|
||||
});
|
||||
|
||||
it('returns false when token cost exceeds budget', () => {
|
||||
// Insert entry with large token count to exceed a tiny budget
|
||||
captureInteraction(store, {
|
||||
sessionId: 's1', workspaceId: 'ws1', systemPrompt: 'p',
|
||||
response: makeResponse({ usage: { inputTokens: 1_000_000, outputTokens: 500_000 } }),
|
||||
turnCount: 1, wasCorrection: false,
|
||||
});
|
||||
|
||||
// 1M input * $3/M + 500K output * $15/M = $3 + $7.5 = $10.50 = 1050 cents
|
||||
// Budget is 1 cent, so way over budget
|
||||
expect(isWithinBudget(store, 1)).toBe(false);
|
||||
});
|
||||
|
||||
it('respects custom cost-per-token rates', () => {
|
||||
captureInteraction(store, {
|
||||
sessionId: 's1', workspaceId: 'ws1', systemPrompt: 'p',
|
||||
response: makeResponse({ usage: { inputTokens: 10_000, outputTokens: 5_000 } }),
|
||||
turnCount: 1, wasCorrection: false,
|
||||
});
|
||||
|
||||
// With very cheap rates: 10K * $0.0001/M + 5K * $0.0001/M ≈ ~0 cents
|
||||
expect(isWithinBudget(store, 1, 0.0001 / 1_000_000, 0.0001 / 1_000_000)).toBe(true);
|
||||
|
||||
// With very expensive rates: 10K * $100/M + 5K * $100/M = $1 + $0.5 = 150 cents
|
||||
expect(isWithinBudget(store, 1, 100 / 1_000_000, 100 / 1_000_000)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
// ── Store stats ────────────────────────────────────────────
|
||||
|
||||
describe('getStats (via store)', () => {
|
||||
it('computes correct aggregate stats', () => {
|
||||
captureInteraction(store, {
|
||||
sessionId: 's1', workspaceId: 'ws1', systemPrompt: 'p',
|
||||
response: makeResponse({ usage: { inputTokens: 100, outputTokens: 50 } }),
|
||||
turnCount: 3, wasCorrection: false,
|
||||
});
|
||||
captureInteraction(store, {
|
||||
sessionId: 's2', workspaceId: 'ws1', systemPrompt: 'p',
|
||||
response: makeResponse({ usage: { inputTokens: 200, outputTokens: 100 } }),
|
||||
turnCount: 5, wasCorrection: true,
|
||||
});
|
||||
|
||||
const stats = store.getStats();
|
||||
expect(stats.total).toBe(2);
|
||||
expect(stats.correctionRate).toBe(0.5); // 1/2
|
||||
expect(stats.avgTurnCount).toBe(4); // (3+5)/2
|
||||
expect(stats.totalInputTokens).toBe(300); // 100+200
|
||||
expect(stats.totalOutputTokens).toBe(150); // 50+100
|
||||
});
|
||||
|
||||
it('returns zeros when empty', () => {
|
||||
const stats = store.getStats();
|
||||
expect(stats.total).toBe(0);
|
||||
expect(stats.correctionRate).toBe(0);
|
||||
expect(stats.avgTurnCount).toBe(0);
|
||||
expect(stats.totalInputTokens).toBe(0);
|
||||
expect(stats.totalOutputTokens).toBe(0);
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user