William Valentin
85 lines
2.3 KiB
TypeScript
85 lines
2.3 KiB
TypeScript
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
|
import { LlamaCppClient } from './llamacpp.js';
|
|
import type { ChatStreamEvent } from '../types.js';
|
|
|
|
describe('LlamaCppClient', () => {
|
|
const mockFetch = vi.fn();
|
|
|
|
beforeEach(() => {
|
|
vi.stubGlobal('fetch', mockFetch);
|
|
});
|
|
|
|
afterEach(() => {
|
|
vi.unstubAllGlobals();
|
|
});
|
|
|
|
it('sends messages and returns response', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: () => Promise.resolve({
|
|
choices: [{ message: { content: 'Hello from llama.cpp!' } }],
|
|
usage: { prompt_tokens: 10, completion_tokens: 5 },
|
|
}),
|
|
});
|
|
|
|
const client = new LlamaCppClient({
|
|
endpoint: 'http://localhost:8080',
|
|
});
|
|
|
|
const response = await client.chat({
|
|
messages: [{ role: 'user', content: 'Hello' }],
|
|
});
|
|
|
|
expect(response.content).toBe('Hello from llama.cpp!');
|
|
expect(response.usage.inputTokens).toBe(10);
|
|
expect(response.usage.outputTokens).toBe(5);
|
|
});
|
|
|
|
it('streams responses via SSE', async () => {
|
|
const chunks = [
|
|
'data: {"choices":[{"delta":{"content":"Hello"}}]}\n\n',
|
|
'data: {"choices":[{"delta":{"content":" world"}}]}\n\n',
|
|
'data: {"choices":[{}],"usage":{"prompt_tokens":5,"completion_tokens":2}}\n\n',
|
|
'data: [DONE]\n\n',
|
|
];
|
|
|
|
const encoder = new TextEncoder();
|
|
let chunkIndex = 0;
|
|
|
|
const mockStream = new ReadableStream({
|
|
pull(controller) {
|
|
if (chunkIndex < chunks.length) {
|
|
controller.enqueue(encoder.encode(chunks[chunkIndex]));
|
|
chunkIndex++;
|
|
} else {
|
|
controller.close();
|
|
}
|
|
},
|
|
});
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: mockStream,
|
|
});
|
|
|
|
const client = new LlamaCppClient({
|
|
endpoint: 'http://localhost:8080',
|
|
});
|
|
|
|
const events: ChatStreamEvent[] = [];
|
|
for await (const event of client.chatStream({
|
|
messages: [{ role: 'user', content: 'Hi' }],
|
|
})) {
|
|
events.push(event);
|
|
}
|
|
|
|
expect(events).toHaveLength(3);
|
|
expect(events[0]).toEqual({ type: 'content', content: 'Hello' });
|
|
expect(events[1]).toEqual({ type: 'content', content: ' world' });
|
|
expect(events[2]).toEqual({
|
|
type: 'done',
|
|
usage: { inputTokens: 5, outputTokens: 2 },
|
|
});
|
|
});
|
|
});
|