feat: SSE parser and buffered chat consumption

This commit is contained in:
zeyu.fz
2026-08-15 18:36:43 +08:00
parent cdc1f10882
commit 95708ccf32
4 changed files with 151 additions and 0 deletions
+41
View File
@@ -0,0 +1,41 @@
/** Buffered consumption of the knowledge chat SSE stream: deltas concatenate into one complete answer. */
import type { ChatStreamChunk } from './api-types.js'
import { KbApiError } from './client.js'
import { parseSseStream } from './sse.js'
export interface ChatResult {
answer: string
requestId?: string
}
/**
* Consume one chat SSE response to completion.
* @param res - the SSE response from KbClient.postSse.
* @returns the concatenated answer and the last seen request id.
*/
export async function consumeChatStream(res: Response): Promise<ChatResult> {
if (!res.body) throw new KbApiError('knowledge chat returned no response body')
let answer = ''
let requestId: string | undefined
for await (const event of parseSseStream(res.body)) {
if (event.data === '[DONE]') break
if (event.event === 'error') {
let message = `knowledge chat stream error: ${event.data}`
try {
const err = JSON.parse(event.data) as { code?: string; message?: string }
if (err.message) message = `knowledge chat stream error${err.code ? ` (${err.code})` : ''}: ${err.message}`
} catch { /* non-JSON error payload: keep the raw data in the message */ }
throw new KbApiError(message)
}
let parsed: ChatStreamChunk
try {
parsed = JSON.parse(event.data) as ChatStreamChunk
} catch { continue } // unparseable keep-alive/comment payloads carry no answer content
if (parsed.request_id) requestId = parsed.request_id
for (const choice of parsed.output?.choices ?? []) {
if (choice.message?.content) answer += choice.message.content
}
}
return { answer, requestId }
}
+51
View File
@@ -0,0 +1,51 @@
/** Minimal SSE parser for the knowledge chat stream: `event:`/`data:` lines, events split on blank lines. */
export interface SseEvent {
event?: string
data: string
}
/**
* Parse one SSE byte stream into events.
* @param body - the response body stream.
* @returns events in stream order; multi-`data:` events join with newlines per the SSE spec.
*/
export async function* parseSseStream(body: ReadableStream<Uint8Array>): AsyncGenerator<SseEvent> {
const decoder = new TextDecoder()
let buffer = ''
let event: string | undefined
let data: string[] = []
const flush = (): SseEvent | undefined => {
if (data.length === 0) return undefined
const out = { event, data: data.join('\n') }
event = undefined
data = []
return out
}
const reader = body.getReader()
while (true) {
const { done, value } = await reader.read()
buffer += done ? '' : decoder.decode(value, { stream: true })
let newline: number
while ((newline = buffer.indexOf('\n')) !== -1) {
const line = buffer.slice(0, newline).replace(/\r$/, '')
buffer = buffer.slice(newline + 1)
if (line === '') {
const out = flush()
if (out) yield out
} else if (line.startsWith('event:')) {
event = line.slice(6).trim()
} else if (line.startsWith('data:')) {
data.push(line.slice(5).trimStart())
}
// comment/id/retry lines are irrelevant to this API and are skipped
}
if (done) {
const out = flush()
if (out) yield out
return
}
}
}
@@ -0,0 +1,30 @@
import { describe, expect, it } from 'vitest'
import { consumeChatStream } from '../src/chat.js'
function sse(text: string): Response {
return new Response(text, { status: 200 })
}
function chunk(content: string, finish = ''): string {
return `data: ${JSON.stringify({ output: { choices: [{ message: { content }, finish_reason: finish }] }, request_id: 'r-1' })}\n\n`
}
describe('consumeChatStream', () => {
it('concatenates delta content across chunks until [DONE]', async () => {
const res = sse(chunk('Hello') + chunk(' world', 'stop') + 'data: [DONE]\n\n')
const out = await consumeChatStream(res)
expect(out.answer).toBe('Hello world')
expect(out.requestId).toBe('r-1')
})
it('ignores step_change progress chunks with empty content', async () => {
const progress = `data: ${JSON.stringify({ output: { choices: [{ message: { content: '', extra: { step_change: 'tool_calling' } }, finish_reason: '' }] } })}\n\n`
const res = sse(progress + chunk('answer', 'stop') + 'data: [DONE]\n\n')
expect((await consumeChatStream(res)).answer).toBe('answer')
})
it('throws on an SSE error event with the server message', async () => {
const res = sse('event: error\ndata: {"code":"Throttling","message":"rate limited"}\n\n')
await expect(consumeChatStream(res)).rejects.toThrow(/Throttling.*rate limited/)
})
})
@@ -0,0 +1,29 @@
import { describe, expect, it } from 'vitest'
import { parseSseStream } from '../src/sse.js'
function streamOf(text: string): ReadableStream<Uint8Array> {
return new Response(text).body as ReadableStream<Uint8Array>
}
async function collect(text: string) {
const events: { event?: string; data: string }[] = []
for await (const e of parseSseStream(streamOf(text))) events.push(e)
return events
}
describe('parseSseStream', () => {
it('yields data events split on blank lines', async () => {
const events = await collect('data: {"a":1}\n\ndata: [DONE]\n\n')
expect(events).toEqual([{ event: undefined, data: '{"a":1}' }, { event: undefined, data: '[DONE]' }])
})
it('carries the event field and parses CRLF lines', async () => {
const events = await collect('event: error\r\ndata: {"message":"boom"}\r\n\r\n')
expect(events[0]).toEqual({ event: 'error', data: '{"message":"boom"}' })
})
it('flushes a final event not terminated by a blank line', async () => {
const events = await collect('data: tail\n')
expect(events).toEqual([{ event: undefined, data: 'tail' }])
})
})