52 lines
1.5 KiB
TypeScript
52 lines
1.5 KiB
TypeScript
import { afterEach, describe, expect, it, vi } from 'vitest';
|
|
|
|
import { streamPlainTextCompletion } from './llmClient';
|
|
|
|
function createSseResponse(body: string) {
|
|
const encoder = new TextEncoder();
|
|
const stream = new ReadableStream<Uint8Array>({
|
|
start(controller) {
|
|
controller.enqueue(encoder.encode(body));
|
|
controller.close();
|
|
},
|
|
});
|
|
|
|
return new Response(stream, {
|
|
headers: {
|
|
'Content-Type': 'text/event-stream; charset=utf-8',
|
|
},
|
|
});
|
|
}
|
|
|
|
describe('llmClient streamPlainTextCompletion', () => {
|
|
afterEach(() => {
|
|
vi.unstubAllGlobals();
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
it('reads OpenAI compatible SSE through the shared stream reader', async () => {
|
|
const onUpdate = vi.fn();
|
|
const fetchMock = vi.fn().mockResolvedValue(
|
|
createSseResponse(
|
|
[
|
|
'data: {"choices":[{"delta":{"content":"溪上"}}]}\r\n\r\n',
|
|
'data: not-json\r\n\r\n',
|
|
'data: {"choices":[{"delta":{"content":"春风"}}]}\r\n\r\n',
|
|
'data: [DONE]\r\n\r\n',
|
|
'data: {"choices":[{"delta":{"content":"不应读取"}}]}\r\n\r\n',
|
|
].join(''),
|
|
),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
|
|
const result = await streamPlainTextCompletion('system', 'user', {
|
|
onUpdate,
|
|
});
|
|
|
|
expect(result).toBe('溪上春风');
|
|
expect(onUpdate).toHaveBeenNthCalledWith(1, '溪上');
|
|
expect(onUpdate).toHaveBeenNthCalledWith(2, '溪上春风');
|
|
expect(onUpdate).toHaveBeenCalledTimes(2);
|
|
});
|
|
});
|