62 lines
1.3 KiB
Go
62 lines
1.3 KiB
Go
package llm
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"io"
|
|
)
|
|
|
|
// sseEvent 是一条 Server-Sent Events 事件。
|
|
type sseEvent struct {
|
|
Name string
|
|
Data []byte
|
|
}
|
|
|
|
// readSSE 逐条解析 SSE 流。兼容 OpenAI 的 "data: {...}" 格式，
|
|
// 也兼容心跳注释行（以 ":" 开头）与自定义 event 名。
|
|
func readSSE(r io.Reader, handle func(sseEvent) error) error {
|
|
br := bufio.NewReaderSize(r, 64*1024)
|
|
var data bytes.Buffer
|
|
var name string
|
|
|
|
flush := func() error {
|
|
if data.Len() == 0 {
|
|
return nil
|
|
}
|
|
ev := sseEvent{Name: name, Data: bytes.TrimSpace(data.Bytes())}
|
|
data.Reset()
|
|
name = ""
|
|
return handle(ev)
|
|
}
|
|
|
|
for {
|
|
line, err := br.ReadBytes('\n')
|
|
if len(line) > 0 {
|
|
line = bytes.TrimRight(line, "\r\n")
|
|
switch {
|
|
case len(line) == 0:
|
|
if ferr := flush(); ferr != nil {
|
|
return ferr
|
|
}
|
|
case bytes.HasPrefix(line, []byte(":")):
|
|
// 注释/心跳，忽略
|
|
case bytes.HasPrefix(line, []byte("data:")):
|
|
v := bytes.TrimPrefix(bytes.TrimPrefix(line, []byte("data:")), []byte(" "))
|
|
if data.Len() > 0 {
|
|
data.WriteByte('\n')
|
|
}
|
|
data.Write(v)
|
|
case bytes.HasPrefix(line, []byte("event:")):
|
|
name = string(bytes.TrimSpace(bytes.TrimPrefix(line, []byte("event:"))))
|
|
}
|
|
}
|
|
if err != nil {
|
|
_ = flush()
|
|
if err == io.EOF {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
}
|