281 lines
10 KiB
TypeScript
281 lines
10 KiB
TypeScript
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);
|
|
});
|
|
});
|
|
});
|