Adapt different stream data prefix, with or without space (#945)
This commit is contained in:
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user