Adapt different stream data prefix, with or without space (#945)

This commit is contained in:
Zhongxian Pan
2025-04-29 21:29:15 +08:00
committed by GitHub
parent 306fbbbe6f
commit 4cccc6c934

View File

@@ -6,13 +6,14 @@ import (
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
"regexp"
utils "github.com/sashabaranov/go-openai/internal" utils "github.com/sashabaranov/go-openai/internal"
) )
var ( var (
headerData = []byte("data: ") headerData = regexp.MustCompile(`^data:\s*`)
errorPrefix = []byte(`data: {"error":`) errorPrefix = regexp.MustCompile(`^data:\s*{"error":`)
) )
type streamable interface { type streamable interface {
@@ -70,12 +71,12 @@ func (stream *streamReader[T]) processLines() ([]byte, error) {
} }
noSpaceLine := bytes.TrimSpace(rawLine) noSpaceLine := bytes.TrimSpace(rawLine)
if bytes.HasPrefix(noSpaceLine, errorPrefix) { if errorPrefix.Match(noSpaceLine) {
hasErrorPrefix = true hasErrorPrefix = true
} }
if !bytes.HasPrefix(noSpaceLine, headerData) || hasErrorPrefix { if !headerData.Match(noSpaceLine) || hasErrorPrefix {
if hasErrorPrefix { if hasErrorPrefix {
noSpaceLine = bytes.TrimPrefix(noSpaceLine, headerData) noSpaceLine = headerData.ReplaceAll(noSpaceLine, nil)
} }
writeErr := stream.errAccumulator.Write(noSpaceLine) writeErr := stream.errAccumulator.Write(noSpaceLine)
if writeErr != nil { if writeErr != nil {
@@ -89,7 +90,7 @@ func (stream *streamReader[T]) processLines() ([]byte, error) {
continue continue
} }
noPrefixLine := bytes.TrimPrefix(noSpaceLine, headerData) noPrefixLine := headerData.ReplaceAll(noSpaceLine, nil)
if string(noPrefixLine) == "[DONE]" { if string(noPrefixLine) == "[DONE]" {
stream.isFinished = true stream.isFinished = true
return nil, io.EOF return nil, io.EOF