100 lines
2.9 KiB
TypeScript
100 lines
2.9 KiB
TypeScript
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
|
import { streamChat, streamGenerate } from './stream.js';
|
|
|
|
const BASE_URL = 'http://localhost:11434';
|
|
|
|
function createNdjsonResponse(chunks: object[]): Response {
|
|
const encoder = new TextEncoder();
|
|
const ndjson = chunks.map(c => JSON.stringify(c)).join('\n') + '\n';
|
|
const stream = new ReadableStream<Uint8Array>({
|
|
start(controller) {
|
|
controller.enqueue(encoder.encode(ndjson));
|
|
controller.close();
|
|
},
|
|
});
|
|
return { ok: true, status: 200, body: stream } as unknown as Response;
|
|
}
|
|
|
|
describe('streamChat', () => {
|
|
beforeEach(() => {
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
it('yields chat stream chunks', async () => {
|
|
const chunks = [
|
|
{ model: 'llama3', message: { role: 'assistant', content: 'Hello' }, done: false },
|
|
{ model: 'llama3', message: { role: 'assistant', content: ' world' }, done: false },
|
|
{ model: 'llama3', message: { role: 'assistant', content: '' }, done: true, eval_count: 10 },
|
|
];
|
|
globalThis.fetch = vi.fn().mockResolvedValue(createNdjsonResponse(chunks));
|
|
|
|
const results = [];
|
|
for await (const chunk of streamChat(BASE_URL, {
|
|
model: 'llama3',
|
|
messages: [{ role: 'user', content: 'Hi' }],
|
|
})) {
|
|
results.push(chunk);
|
|
}
|
|
|
|
expect(results).toHaveLength(3);
|
|
expect(results[0].message.content).toBe('Hello');
|
|
expect(results[2].done).toBe(true);
|
|
expect(results[2].eval_count).toBe(10);
|
|
});
|
|
|
|
it('throws on non-ok response', async () => {
|
|
globalThis.fetch = vi.fn().mockResolvedValue({
|
|
ok: false,
|
|
status: 500,
|
|
text: () => Promise.resolve('internal error'),
|
|
});
|
|
|
|
const gen = streamChat(BASE_URL, {
|
|
model: 'llama3',
|
|
messages: [{ role: 'user', content: 'Hi' }],
|
|
});
|
|
await expect(gen.next()).rejects.toThrow('Ollama chat failed (500)');
|
|
});
|
|
|
|
it('throws when response has no body', async () => {
|
|
globalThis.fetch = vi.fn().mockResolvedValue({
|
|
ok: true,
|
|
status: 200,
|
|
body: null,
|
|
});
|
|
|
|
const gen = streamChat(BASE_URL, {
|
|
model: 'llama3',
|
|
messages: [{ role: 'user', content: 'Hi' }],
|
|
});
|
|
await expect(gen.next()).rejects.toThrow('No response body');
|
|
});
|
|
});
|
|
|
|
describe('streamGenerate', () => {
|
|
beforeEach(() => {
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
it('yields generate stream chunks', async () => {
|
|
const chunks = [
|
|
{ model: 'llama3', response: 'Hello', done: false },
|
|
{ model: 'llama3', response: ' world', done: false },
|
|
{ model: 'llama3', response: '', done: true, eval_count: 5 },
|
|
];
|
|
globalThis.fetch = vi.fn().mockResolvedValue(createNdjsonResponse(chunks));
|
|
|
|
const results = [];
|
|
for await (const chunk of streamGenerate(BASE_URL, {
|
|
model: 'llama3',
|
|
prompt: 'Say hello',
|
|
})) {
|
|
results.push(chunk);
|
|
}
|
|
|
|
expect(results).toHaveLength(3);
|
|
expect(results[0].response).toBe('Hello');
|
|
expect(results[2].done).toBe(true);
|
|
});
|
|
});
|