fix(model):openai格式对齐new-api
This commit is contained in:
qian committed
1 parent
9334671ea7
commit
ce9a20366c
1 file changed
+57
-22
@@ -30,13 +30,22 @@ const (
|
||||
// DataHandlerFunc data 处理回调
|
||||
type DataHandlerFunc func(data string, sr *StreamResult)
|
||||
|
||||
// StreamScannerHandler 三协程 SSE 流扫描器
|
||||
// Scanner 协程：读取上游行，过滤 data: 行，推入 dataChan
|
||||
// StreamScannerHandler 三协程 SSE 流扫描器（对齐 new-api 架构）
|
||||
//
|
||||
// Scanner 协程：读取上游行，检测 [DONE] 终止标记并立即退出（不转发 [DONE] 到 dataChan）
|
||||
// DataHandler 协程：从 dataChan 读取，调用 handler 回调，持有 writeMutex
|
||||
// Ping 协程：定期写 SSE 保活注释，持有 writeMutex
|
||||
//
|
||||
// 超时策略：per-chunk 空闲超时（参考 new-api），每收到一个有效 chunk 就重置计时器。
|
||||
// 只要上游持续发送数据，流可以无限延续，适用于长时间思考的模型。
|
||||
// 超时策略：per-chunk 空闲超时，scanner.Scan() 每返回一行就通过 resetChan 通知主协程重置 timer。
|
||||
// scanner.Scan() 阻塞期间不重置，超时从最后收到的字节开始算。
|
||||
// 对齐 new-api 的 ticker.Reset(streamingTimeout) 行为，但用 Timer+resetChan 避免 ticker 跨协程操作。
|
||||
//
|
||||
// 终止流程（对齐 new-api）：
|
||||
// 1. Scanner 检测到 [DONE] → return → close(dataChan) → safeSendBool(stopChan)
|
||||
// 2. DataHandler range 退出 → safeSendBool(stopChan)
|
||||
// 3. 主 goroutine 收到 stopChan → 进入 shutdown
|
||||
// 4. shutdown: safeSendBool(stopChan) 通知剩余协程 → wg.Wait（最多 5s） → resp.Body.Close()
|
||||
// 5. [DONE] 由 handler 层（如 openai/stream.go）负责写入客户端
|
||||
func StreamScannerHandler(
|
||||
ctx context.Context,
|
||||
resp *http.Response,
|
||||
@@ -46,24 +55,38 @@ func StreamScannerHandler(
|
||||
) {
|
||||
stopChan := make(chan bool, stopChanSize)
|
||||
dataChan := make(chan string, dataChanSize)
|
||||
resetChan := make(chan struct{}, resetChanSize) // 收到有效数据时通知主 goroutine 重置超时
|
||||
resetChan := make(chan struct{}, resetChanSize) // scanner 每读一行通知主协程重置超时
|
||||
var writeMutex sync.Mutex
|
||||
var wg sync.WaitGroup
|
||||
streamTimeout := getStreamTimeout(info)
|
||||
|
||||
// 确保 resp.Body 最终被关闭（对齐 new-api），解锁被 scanner.Scan() 阻塞的协程
|
||||
defer func() {
|
||||
if resp.Body != nil {
|
||||
resp.Body.Close()
|
||||
}
|
||||
}()
|
||||
|
||||
SetEventStreamHeaders(writer)
|
||||
|
||||
// Scanner 协程
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer close(dataChan)
|
||||
defer func() {
|
||||
close(dataChan)
|
||||
wg.Done()
|
||||
if r := recover(); r != nil {
|
||||
info.StreamStatus.SetEndReason(common.StreamEndReasonError, fmt.Errorf("scanner panic: %v", r))
|
||||
}
|
||||
safeSendBool(stopChan)
|
||||
}()
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
buf := make([]byte, 0, scannerInitialBuf)
|
||||
scanner.Buffer(buf, scannerMaxBuf)
|
||||
|
||||
for scanner.Scan() {
|
||||
// 检查停止信号
|
||||
select {
|
||||
case <-stopChan:
|
||||
return
|
||||
@@ -73,13 +96,22 @@ func StreamScannerHandler(
|
||||
default:
|
||||
}
|
||||
|
||||
// 每读一行就通知主协程重置空闲超时（对齐 new-api ticker.Reset）
|
||||
select {
|
||||
case resetChan <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
|
||||
line := scanner.Text()
|
||||
if len(line) < minLineLength {
|
||||
continue
|
||||
}
|
||||
|
||||
// 兼容两种 [DONE] 格式：
|
||||
// 1. 裸 "[DONE]"（少数上游）
|
||||
// 2. 标准 SSE "data: [DONE]"（主流上游如 OpenAI/DeepSeek）
|
||||
// 检测到后立即退出 scanner，由 handler 层负责向客户端写入 [DONE]。
|
||||
if line == "[DONE]" {
|
||||
safeSendBool(stopChan)
|
||||
info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil)
|
||||
return
|
||||
}
|
||||
@@ -93,14 +125,13 @@ func StreamScannerHandler(
|
||||
continue
|
||||
}
|
||||
|
||||
if info.StreamStatus.GetEndReason() == "" {
|
||||
info.SetFirstResponseTime()
|
||||
if data == "[DONE]" {
|
||||
info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil)
|
||||
return
|
||||
}
|
||||
|
||||
// 通知主 goroutine 重置空闲超时
|
||||
select {
|
||||
case resetChan <- struct{}{}:
|
||||
default:
|
||||
if info.StreamStatus.GetEndReason() == "" {
|
||||
info.SetFirstResponseTime()
|
||||
}
|
||||
|
||||
select {
|
||||
@@ -115,10 +146,8 @@ func StreamScannerHandler(
|
||||
|
||||
if err := scanner.Err(); err != nil && err != io.EOF && ctx.Err() == nil {
|
||||
info.StreamStatus.SetEndReason(common.StreamEndReasonScannerErr, err)
|
||||
safeSendBool(stopChan)
|
||||
} else if info.StreamStatus.GetEndReason() == "" {
|
||||
info.StreamStatus.SetEndReason(common.StreamEndReasonEOF, nil)
|
||||
safeSendBool(stopChan)
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -126,14 +155,19 @@ func StreamScannerHandler(
|
||||
sr := NewStreamResult(info.StreamStatus)
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer func() {
|
||||
wg.Done()
|
||||
if r := recover(); r != nil {
|
||||
info.StreamStatus.SetEndReason(common.StreamEndReasonError, fmt.Errorf("handler panic: %v", r))
|
||||
}
|
||||
safeSendBool(stopChan)
|
||||
}()
|
||||
for data := range dataChan {
|
||||
sr.reset()
|
||||
writeMutex.Lock()
|
||||
handler(data, sr)
|
||||
writeMutex.Unlock()
|
||||
if sr.IsStopped() {
|
||||
safeSendBool(stopChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -168,8 +202,7 @@ func StreamScannerHandler(
|
||||
}
|
||||
}()
|
||||
|
||||
// 主 goroutine 阻塞等待
|
||||
// 使用 Timer 而非 Ticker，每次收到有效数据时 Reset，实现 per-chunk 空闲超时
|
||||
// 主 goroutine 阻塞等待（对齐 new-api：单次 select，用 Timer 实现空闲超时）
|
||||
timeoutTimer := time.NewTimer(streamTimeout)
|
||||
defer timeoutTimer.Stop()
|
||||
|
||||
@@ -185,7 +218,7 @@ func StreamScannerHandler(
|
||||
info.StreamStatus.SetEndReason(common.StreamEndReasonClientGone, ctx.Err())
|
||||
goto shutdown
|
||||
case <-resetChan:
|
||||
// 收到有效数据，重置空闲超时
|
||||
// scanner 每读一行都通知重置（对齐 new-api ticker.Reset）
|
||||
if !timeoutTimer.Stop() {
|
||||
select {
|
||||
case <-timeoutTimer.C:
|
||||
@@ -197,7 +230,9 @@ func StreamScannerHandler(
|
||||
}
|
||||
|
||||
shutdown:
|
||||
// 等待所有协程退出
|
||||
// 通知所有协程停止（对齐 new-api：先 signal 再 wait）
|
||||
safeSendBool(stopChan)
|
||||
|
||||
waitDone := make(chan struct{})
|
||||
go func() {
|
||||
wg.Wait()
|
||||
|
||||
Reference in new issue
Block a user