369 lines
11 KiB
TypeScript
369 lines
11 KiB
TypeScript
import { OllamaClient } from '../src/ollama-client';
|
|
import { OllamaMessage, OllamaTool } from '../src/types';
|
|
|
|
describe('OllamaClient', () => {
|
|
let client: OllamaClient;
|
|
let mockFetch: jest.Mock;
|
|
|
|
const mockMessages: OllamaMessage[] = [
|
|
{ role: 'system', content: 'You are helpful.' },
|
|
{ role: 'user', content: 'Hello' },
|
|
];
|
|
|
|
const mockTools: OllamaTool[] = [
|
|
{
|
|
type: 'function',
|
|
function: {
|
|
name: 'test_tool',
|
|
description: 'A test tool',
|
|
parameters: {
|
|
type: 'object',
|
|
properties: { input: { type: 'string' } },
|
|
required: ['input'],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
|
|
beforeEach(() => {
|
|
mockFetch = jest.fn();
|
|
client = new OllamaClient('http://localhost:11434', 'llama3', mockFetch);
|
|
});
|
|
|
|
afterEach(() => {
|
|
jest.clearAllMocks();
|
|
client.cancelStream();
|
|
});
|
|
|
|
describe('chat (non-streaming)', () => {
|
|
it('should send a non-streaming request and return the response', async () => {
|
|
const mockResponse = {
|
|
ok: true,
|
|
json: () => Promise.resolve({ message: { content: 'Hello back!' } }),
|
|
};
|
|
mockFetch.mockResolvedValue(mockResponse);
|
|
|
|
const result = await client.chat(mockMessages, mockTools);
|
|
|
|
expect(result.content).toBe('Hello back!');
|
|
expect(mockFetch).toHaveBeenCalledWith(
|
|
'http://localhost:11434/api/chat',
|
|
expect.objectContaining({
|
|
method: 'POST',
|
|
body: JSON.stringify({
|
|
model: 'llama3',
|
|
messages: mockMessages,
|
|
tools: mockTools,
|
|
stream: false,
|
|
}),
|
|
})
|
|
);
|
|
});
|
|
|
|
it('should throw on non-OK response', async () => {
|
|
mockFetch.mockResolvedValue({ ok: false, status: 500 });
|
|
|
|
await expect(client.chat(mockMessages, mockTools)).rejects.toThrow('Ollama API error: 500');
|
|
});
|
|
|
|
it('should handle missing message content gracefully', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: () => Promise.resolve({}),
|
|
});
|
|
|
|
const result = await client.chat(mockMessages, mockTools);
|
|
expect(result.content).toBe('');
|
|
expect(result.tool_calls).toEqual([]);
|
|
});
|
|
|
|
it('should include abort signal in fetch options', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: () => Promise.resolve({ message: { content: 'ok' } }),
|
|
});
|
|
|
|
await client.chat(mockMessages, mockTools);
|
|
|
|
const fetchOptions = mockFetch.mock.calls[0][1];
|
|
expect(fetchOptions.signal).toBeInstanceOf(AbortSignal);
|
|
});
|
|
|
|
it('should forward tool_calls from response when present', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: () =>
|
|
Promise.resolve({
|
|
message: {
|
|
content: 'result',
|
|
tool_calls: [{ function: { name: 'create_file', arguments: '{}' } }],
|
|
},
|
|
}),
|
|
});
|
|
|
|
const result = await client.chat(mockMessages, mockTools);
|
|
expect(result.tool_calls).toEqual([{ function: { name: 'create_file', arguments: '{}' } }]);
|
|
});
|
|
});
|
|
|
|
describe('streamChat', () => {
|
|
function createMockReader(data: string) {
|
|
const encoder = new TextEncoder();
|
|
const encoded = encoder.encode(data);
|
|
let called = false;
|
|
return {
|
|
read: () => {
|
|
if (!called) {
|
|
called = true;
|
|
return Promise.resolve({ done: false, value: encoded });
|
|
}
|
|
return Promise.resolve({ done: true, value: new Uint8Array(0) });
|
|
},
|
|
releaseLock: jest.fn(),
|
|
};
|
|
}
|
|
|
|
it('should send a streaming request and yield chunks', async () => {
|
|
const streamData = [
|
|
JSON.stringify({ message: { content: 'He' } }),
|
|
JSON.stringify({ message: { content: 'llo' } }),
|
|
JSON.stringify({ message: { content: '!' } }),
|
|
'',
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
const chunks: string[] = [];
|
|
|
|
for await (const chunk of stream) {
|
|
chunks.push(chunk.content);
|
|
}
|
|
|
|
expect(chunks).toEqual(['He', 'llo', '!']);
|
|
expect(mockReader.releaseLock).toHaveBeenCalled();
|
|
});
|
|
|
|
it('should skip malformed JSON chunks and log a warning', async () => {
|
|
const streamData = [
|
|
JSON.stringify({ message: { content: 'valid' } }),
|
|
'this is not json',
|
|
JSON.stringify({ message: { content: 'also valid' } }),
|
|
'',
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
const chunks: string[] = [];
|
|
|
|
for await (const chunk of stream) {
|
|
chunks.push(chunk.content);
|
|
}
|
|
|
|
expect(chunks).toEqual(['valid', 'also valid']);
|
|
expect(consoleWarnSpy).toHaveBeenCalledWith(
|
|
expect.stringContaining('Skipped malformed chunk'),
|
|
expect.stringContaining('is not valid JSON')
|
|
);
|
|
consoleWarnSpy.mockRestore();
|
|
});
|
|
|
|
it('should throw when too many chunks are malformed', async () => {
|
|
const streamData = Array(51).fill('invalid json').join('\n') + '\n';
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
|
|
await expect(
|
|
(async () => {
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
})()
|
|
).rejects.toThrow(/malformed/);
|
|
});
|
|
|
|
it('should throw on non-OK response', async () => {
|
|
mockFetch.mockResolvedValue({ ok: false, status: 404 });
|
|
|
|
await expect(client.streamChat(mockMessages, mockTools)).rejects.toThrow(
|
|
'Ollama API error: 404'
|
|
);
|
|
});
|
|
|
|
it('should throw when response has no body', async () => {
|
|
mockFetch.mockResolvedValue({ ok: true, body: undefined });
|
|
|
|
await expect(client.streamChat(mockMessages, mockTools)).rejects.toThrow('No response body');
|
|
});
|
|
|
|
it('should throw on invalid content type', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: {
|
|
getReader: () => ({
|
|
read: () => Promise.resolve({ done: true, value: new Uint8Array(0) }),
|
|
}),
|
|
},
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'text/html' : null),
|
|
},
|
|
});
|
|
|
|
await expect(client.streamChat(mockMessages, mockTools)).rejects.toThrow(
|
|
'Invalid response format'
|
|
);
|
|
});
|
|
|
|
it('should propagate Ollama error messages from the stream', async () => {
|
|
const streamData = JSON.stringify({ message: { error: 'model not found' } }) + '\n';
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
|
|
await expect(
|
|
(async () => {
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
})()
|
|
).rejects.toThrow('Ollama error: model not found');
|
|
});
|
|
|
|
it('should yield tool_calls when present in streamed response', async () => {
|
|
const streamData = [
|
|
JSON.stringify({
|
|
message: {
|
|
content: '',
|
|
tool_calls: [{ function: { name: 'create_file', arguments: '{"path":"a.md"}' } }],
|
|
},
|
|
}),
|
|
'',
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
let lastChunk: any;
|
|
|
|
for await (const chunk of stream) {
|
|
lastChunk = chunk;
|
|
}
|
|
|
|
expect(lastChunk.tool_calls).toEqual([
|
|
{ function: { name: 'create_file', arguments: '{"path":"a.md"}' } },
|
|
]);
|
|
});
|
|
|
|
it('should default tool_calls to empty array when not present', async () => {
|
|
const streamData = JSON.stringify({ message: { content: 'hello' } }) + '\n';
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
let lastChunk: any;
|
|
|
|
for await (const chunk of stream) {
|
|
lastChunk = chunk;
|
|
}
|
|
|
|
expect(lastChunk.tool_calls).toEqual([]);
|
|
});
|
|
|
|
it('should send correct request body with stream:true', async () => {
|
|
const streamData = JSON.stringify({ message: { content: 'ok' } }) + '\n';
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
|
|
expect(mockFetch).toHaveBeenCalledWith(
|
|
'http://localhost:11434/api/chat',
|
|
expect.objectContaining({
|
|
method: 'POST',
|
|
body: JSON.stringify({
|
|
model: 'llama3',
|
|
messages: mockMessages,
|
|
tools: mockTools,
|
|
stream: true,
|
|
}),
|
|
signal: expect.any(AbortSignal),
|
|
})
|
|
);
|
|
});
|
|
});
|
|
|
|
describe('cancelStream', () => {
|
|
it('should abort the current request', () => {
|
|
client.cancelStream();
|
|
expect(client['abortController']).toBeNull();
|
|
});
|
|
|
|
it('should handle cancel when no active stream', () => {
|
|
expect(() => client.cancelStream()).not.toThrow();
|
|
expect(client['abortController']).toBeNull();
|
|
});
|
|
});
|
|
});
|