Skip to content

Commit 756dc15

Browse files
committed
Adapt different stream data prefix, with or without space
1 parent 85f578b commit 756dc15

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

stream_reader.go

+7-6
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,14 @@ import (
66
"fmt"
77
"io"
88
"net/http"
9+
"regexp"
910

1011
utils "github.com/sashabaranov/go-openai/internal"
1112
)
1213

1314
var (
14-
headerData = []byte("data: ")
15-
errorPrefix = []byte(`data: {"error":`)
15+
headerData = regexp.MustCompile(`^data:\s*`)
16+
errorPrefix = regexp.MustCompile(`^data:\s*{"error":`)
1617
)
1718

1819
type streamable interface {
@@ -70,12 +71,12 @@ func (stream *streamReader[T]) processLines() ([]byte, error) {
7071
}
7172

7273
noSpaceLine := bytes.TrimSpace(rawLine)
73-
if bytes.HasPrefix(noSpaceLine, errorPrefix) {
74+
if errorPrefix.Match(noSpaceLine) {
7475
hasErrorPrefix = true
7576
}
76-
if !bytes.HasPrefix(noSpaceLine, headerData) || hasErrorPrefix {
77+
if !headerData.Match(noSpaceLine) || hasErrorPrefix {
7778
if hasErrorPrefix {
78-
noSpaceLine = bytes.TrimPrefix(noSpaceLine, headerData)
79+
noSpaceLine = headerData.ReplaceAll(noSpaceLine, nil)
7980
}
8081
writeErr := stream.errAccumulator.Write(noSpaceLine)
8182
if writeErr != nil {
@@ -89,7 +90,7 @@ func (stream *streamReader[T]) processLines() ([]byte, error) {
8990
continue
9091
}
9192

92-
noPrefixLine := bytes.TrimPrefix(noSpaceLine, headerData)
93+
noPrefixLine := headerData.ReplaceAll(noSpaceLine, nil)
9394
if string(noPrefixLine) == "[DONE]" {
9495
stream.isFinished = true
9596
return nil, io.EOF

0 commit comments

Comments
 (0)