Test Utilities

pnpm add -D @cogitator-ai/test-utils

import { Cogitator, Agent, tool } from '@cogitator-ai/core';
import { MockLLMBackend, createToolCall } from '@cogitator-ai/test-utils';
import { z } from 'zod';

const getWeather = tool({
  name: 'get_weather',
  description: 'Get the weather for a city',
  parameters: z.object({ city: z.string() }),
  execute: async ({ city }) => ({ city, temperature: 22 }),
});

const mock = new MockLLMBackend().setResponses([
  { toolCalls: [createToolCall('get_weather', { city: 'Tokyo' })] },
  { content: 'It is 22°C in Tokyo.' },
]);

const cog = new Cogitator({ llm: { backends: { mock } } });
const agent = new Agent({
  name: 'weather',
  model: 'mock/test-model',
  instructions: 'You are a weather assistant.',
  tools: [getWeather],
});

const result = await cog.run(agent, { input: 'Weather in Tokyo?' });

expect(result.output).toBe('It is 22°C in Tokyo.');
expect(result.toolCalls.map((call) => call.name)).toEqual(['get_weather']);
expect(mock.getCallCount()).toBe(2);
expect(mock.getLastCall()?.model).toBe('test-model');

await cog.close();

interface MockChatResponse extends Partial<ChatResponse> {
  content?: string; // default ''
  toolCalls?: ToolCall[];
  finishReason?: 'stop' | 'tool_calls' | 'length' | 'error'; // default: 'tool_calls' with toolCalls, else 'stop'
  error?: Error; // thrown instead of answering
  delay?: number; // ms, overrides setDefaultDelay
}

const mock = new MockLLMBackend().setResponse({ error: new Error('provider down') });
const cog = new Cogitator({ llm: { backends: { mock }, retry: false } });

await expect(cog.run(agent, { input: 'hi' })).rejects.toThrow('provider down');

const mock = new MockLLMBackend().setStreamChunks([
  { content: 'Once ' },
  { content: 'upon a time' },
]);
const cog = new Cogitator({ llm: { backends: { mock } } });

const tokens: string[] = [];
const result = await cog.run(agent, {
  input: 'Tell me a story',
  stream: true,
  onToken: (token) => tokens.push(token),
});

expect(tokens).toEqual(['Once ', 'upon a time']);
expect(result.output).toBe('Once upon a time');

mock.setMultipleStreamChunks([
  [{ toolCalls: [{ id: 'call_1', name: 'get_weather', arguments: { city: 'Oslo' } }] }],
  [{ content: 'It is cold in Oslo.' }],
]);

expect(mock.getCalls()[0].tools?.map((t) => t.name)).toEqual(['get_weather']);
expect(
  mock.wasCalledWith((req) => req.messages[0].content === 'You are a weather assistant.')
).toBe(true);

import type { ChatRequest, ChatResponse, LLMBackend } from '@cogitator-ai/types';

const requests: ChatRequest[] = [];
const backend: LLMBackend = {
  provider: 'openai',
  async chat(request): Promise<ChatResponse> {
    requests.push({ ...request, messages: [...request.messages] });
    const usage = { inputTokens: 1, outputTokens: 1, totalTokens: 2 };
    const hasToolResult = request.messages.some((m) => m.role === 'tool');
    return hasToolResult
      ? { id: 'r2', content: 'done', finishReason: 'stop', usage }
      : {
          id: 'r1',
          content: '',
          toolCalls: [{ id: 'call_1', name: 'get_weather', arguments: { city: 'Paris' } }],
          finishReason: 'tool_calls',
          usage,
        };
  },
  async *chatStream() {
    yield { id: 's', delta: {}, finishReason: 'stop' };
  },
};

const cog = new Cogitator({ llm: { backends: { mock: backend } } });

import { MockMemoryAdapter, MockLLMBackend } from '@cogitator-ai/test-utils';

const memory = new MockMemoryAdapter();
await memory.connect();

const mock = new MockLLMBackend().setResponse({ content: 'Hi there' });
const cog = new Cogitator({ llm: { backends: { mock } } });
cog.memory = memory;

await cog.run(agent, { input: 'Hello', threadId: 'thread-1' });
await cog.run(agent, { input: 'Again', threadId: 'thread-1' });

expect(memory.getCallsFor('createThread')).toHaveLength(1);
expect(memory.getCallsFor('addEntry').length).toBeGreaterThan(0);

import { Agent } from '@cogitator-ai/core';
import { createTestAgentConfig, createTestAgent } from '@cogitator-ai/test-utils';

const agent = new Agent(createTestAgentConfig({ tools: [getWeather] }));

const call = createToolCall('get_weather', { city: 'Tokyo' });
const history = [
  createUserMessage('Weather in Tokyo?'),
  createToolMessage(call.id, JSON.stringify({ temperature: 22 }), 'get_weather'),
];
const result = createMockRunResult('It is 22°C.', { toolCalls: [call], messages: history });

const calculator = createCalculatorTool();
await expect(calculator.execute({ expression: '6 * 7' }, context)).resolves.toEqual({ result: 42 });

const mock = new MockLLMBackend().setStreamChunks([{ content: 'Hel' }, { content: 'lo' }]);
const text = await collectStreamContent(mock.chatStream({ model: 'm', messages: [] }));
expect(text).toBe('Hello');

const job = startBackgroundJob();
await eventually(() => expect(job.done).toBe(true), { timeout: 2000 });