feat: add streaming support to LlamaCppClient
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
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();
|
||||
@@ -33,4 +34,51 @@ describe('LlamaCppClient', () => {
|
||||
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 },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user