mirror of
https://github.com/modelstudioai/cli.git
synced 2026-09-14 19:49:23 +08:00
feat: SSE parser and buffered chat consumption
This commit is contained in:
@@ -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 }
|
||||
}
|
||||
@@ -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' }])
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user