fix: chat stream returns an error response with a 'data: ' prefix (#396)
* fix: chat stream resp has 'data: ' prefix * fix: lint error * fix: lint error * fix: lint error
This commit is contained in:
@@ -178,6 +178,45 @@ func TestCreateChatCompletionStreamError(t *testing.T) {
|
|||||||
t.Logf("%+v\n", apiErr)
|
t.Logf("%+v\n", apiErr)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestCreateChatCompletionStreamErrorWithDataPrefix(t *testing.T) {
|
||||||
|
client, server, teardown := setupOpenAITestServer()
|
||||||
|
defer teardown()
|
||||||
|
server.RegisterHandler("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
|
||||||
|
// Send test responses
|
||||||
|
//nolint:lll
|
||||||
|
dataBytes := []byte(`data: {"error":{"message":"The server had an error while processing your request. Sorry about that!", "type":"server_ error", "param":null,"code":null}}`)
|
||||||
|
dataBytes = append(dataBytes, []byte("\n\ndata: [DONE]\n\n")...)
|
||||||
|
|
||||||
|
_, err := w.Write(dataBytes)
|
||||||
|
checks.NoError(t, err, "Write error")
|
||||||
|
})
|
||||||
|
|
||||||
|
stream, err := client.CreateChatCompletionStream(context.Background(), ChatCompletionRequest{
|
||||||
|
MaxTokens: 5,
|
||||||
|
Model: GPT3Dot5Turbo,
|
||||||
|
Messages: []ChatCompletionMessage{
|
||||||
|
{
|
||||||
|
Role: ChatMessageRoleUser,
|
||||||
|
Content: "Hello!",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Stream: true,
|
||||||
|
})
|
||||||
|
checks.NoError(t, err, "CreateCompletionStream returned error")
|
||||||
|
defer stream.Close()
|
||||||
|
|
||||||
|
_, streamErr := stream.Recv()
|
||||||
|
checks.HasError(t, streamErr, "stream.Recv() did not return error")
|
||||||
|
|
||||||
|
var apiErr *APIError
|
||||||
|
if !errors.As(streamErr, &apiErr) {
|
||||||
|
t.Errorf("stream.Recv() did not return APIError")
|
||||||
|
}
|
||||||
|
t.Logf("%+v\n", apiErr)
|
||||||
|
}
|
||||||
|
|
||||||
func TestCreateChatCompletionStreamRateLimitError(t *testing.T) {
|
func TestCreateChatCompletionStreamRateLimitError(t *testing.T) {
|
||||||
client, server, teardown := setupOpenAITestServer()
|
client, server, teardown := setupOpenAITestServer()
|
||||||
defer teardown()
|
defer teardown()
|
||||||
|
|||||||
@@ -10,6 +10,11 @@ import (
|
|||||||
utils "github.com/sashabaranov/go-openai/internal"
|
utils "github.com/sashabaranov/go-openai/internal"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
headerData = []byte("data: ")
|
||||||
|
errorPrefix = []byte(`data: {"error":`)
|
||||||
|
)
|
||||||
|
|
||||||
type streamable interface {
|
type streamable interface {
|
||||||
ChatCompletionStreamResponse | CompletionResponse
|
ChatCompletionStreamResponse | CompletionResponse
|
||||||
}
|
}
|
||||||
@@ -34,12 +39,16 @@ func (stream *streamReader[T]) Recv() (response T, err error) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
//nolint:gocognit
|
||||||
func (stream *streamReader[T]) processLines() (T, error) {
|
func (stream *streamReader[T]) processLines() (T, error) {
|
||||||
var emptyMessagesCount uint
|
var (
|
||||||
|
emptyMessagesCount uint
|
||||||
|
hasErrorPrefix bool
|
||||||
|
)
|
||||||
|
|
||||||
for {
|
for {
|
||||||
rawLine, readErr := stream.reader.ReadBytes('\n')
|
rawLine, readErr := stream.reader.ReadBytes('\n')
|
||||||
if readErr != nil {
|
if readErr != nil || hasErrorPrefix {
|
||||||
respErr := stream.unmarshalError()
|
respErr := stream.unmarshalError()
|
||||||
if respErr != nil {
|
if respErr != nil {
|
||||||
return *new(T), fmt.Errorf("error, %w", respErr.Error)
|
return *new(T), fmt.Errorf("error, %w", respErr.Error)
|
||||||
@@ -47,9 +56,14 @@ func (stream *streamReader[T]) processLines() (T, error) {
|
|||||||
return *new(T), readErr
|
return *new(T), readErr
|
||||||
}
|
}
|
||||||
|
|
||||||
var headerData = []byte("data: ")
|
|
||||||
noSpaceLine := bytes.TrimSpace(rawLine)
|
noSpaceLine := bytes.TrimSpace(rawLine)
|
||||||
if !bytes.HasPrefix(noSpaceLine, headerData) {
|
if bytes.HasPrefix(noSpaceLine, errorPrefix) {
|
||||||
|
hasErrorPrefix = true
|
||||||
|
}
|
||||||
|
if !bytes.HasPrefix(noSpaceLine, headerData) || hasErrorPrefix {
|
||||||
|
if hasErrorPrefix {
|
||||||
|
noSpaceLine = bytes.TrimPrefix(noSpaceLine, headerData)
|
||||||
|
}
|
||||||
writeErr := stream.errAccumulator.Write(noSpaceLine)
|
writeErr := stream.errAccumulator.Write(noSpaceLine)
|
||||||
if writeErr != nil {
|
if writeErr != nil {
|
||||||
return *new(T), writeErr
|
return *new(T), writeErr
|
||||||
|
|||||||
Reference in New Issue
Block a user