forked from sashabaranov/go-openai
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix: stream return EOF when openai return error (sashabaranov#184)
* fix: stream return EOF when openai return error * perf: add error accumulator * fix: golangci-lint * fix: unmarshal error possibly null * fix: error accumulator * test: error accumulator use interface and add test code * test: error accumulator add test code * refactor: use stream reader to re-use stream code * refactor: stream reader use generics
- Loading branch information
1 parent
aa149c1
commit a5a945a
Showing
8 changed files
with
372 additions
and
107 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
package openai | ||
|
||
import ( | ||
"bytes" | ||
"fmt" | ||
"io" | ||
) | ||
|
||
type errorAccumulator interface { | ||
write(p []byte) error | ||
unmarshalError() (*ErrorResponse, error) | ||
} | ||
|
||
type errorBuffer interface { | ||
io.Writer | ||
Len() int | ||
Bytes() []byte | ||
} | ||
|
||
type errorAccumulate struct { | ||
buffer errorBuffer | ||
unmarshaler unmarshaler | ||
} | ||
|
||
func newErrorAccumulator() errorAccumulator { | ||
return &errorAccumulate{ | ||
buffer: &bytes.Buffer{}, | ||
unmarshaler: &jsonUnmarshaler{}, | ||
} | ||
} | ||
|
||
func (e *errorAccumulate) write(p []byte) error { | ||
_, err := e.buffer.Write(p) | ||
if err != nil { | ||
return fmt.Errorf("error accumulator write error, %w", err) | ||
} | ||
return nil | ||
} | ||
|
||
func (e *errorAccumulate) unmarshalError() (*ErrorResponse, error) { | ||
var err error | ||
if e.buffer.Len() > 0 { | ||
var errRes ErrorResponse | ||
err = e.unmarshaler.unmarshal(e.buffer.Bytes(), &errRes) | ||
if err != nil { | ||
return nil, err | ||
} | ||
return &errRes, nil | ||
} | ||
return nil, err | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,90 @@ | ||
package openai //nolint:testpackage // testing private field | ||
|
||
import ( | ||
"bytes" | ||
"context" | ||
"errors" | ||
"testing" | ||
|
||
"github.com/sashabaranov/go-openai/internal/test" | ||
) | ||
|
||
var ( | ||
errTestUnmarshalerFailed = errors.New("test unmarshaler failed") | ||
errTestErrorAccumulatorWriteFailed = errors.New("test error accumulator failed") | ||
) | ||
|
||
type ( | ||
failingUnMarshaller struct{} | ||
failingErrorBuffer struct{} | ||
) | ||
|
||
func (b *failingErrorBuffer) Write(_ []byte) (n int, err error) { | ||
return 0, errTestErrorAccumulatorWriteFailed | ||
} | ||
|
||
func (b *failingErrorBuffer) Len() int { | ||
return 0 | ||
} | ||
|
||
func (b *failingErrorBuffer) Bytes() []byte { | ||
return []byte{} | ||
} | ||
|
||
func (*failingUnMarshaller) unmarshal(_ []byte, _ any) error { | ||
return errTestUnmarshalerFailed | ||
} | ||
|
||
func TestErrorAccumulatorReturnsUnmarshalerErrors(t *testing.T) { | ||
accumulator := &errorAccumulate{ | ||
buffer: &bytes.Buffer{}, | ||
unmarshaler: &failingUnMarshaller{}, | ||
} | ||
|
||
err := accumulator.write([]byte("{")) | ||
if err != nil { | ||
t.Fatalf("%+v", err) | ||
} | ||
_, err = accumulator.unmarshalError() | ||
if !errors.Is(err, errTestUnmarshalerFailed) { | ||
t.Fatalf("Did not return error when unmarshaler failed: %v", err) | ||
} | ||
} | ||
|
||
func TestErrorByteWriteErrors(t *testing.T) { | ||
accumulator := &errorAccumulate{ | ||
buffer: &failingErrorBuffer{}, | ||
unmarshaler: &jsonUnmarshaler{}, | ||
} | ||
err := accumulator.write([]byte("{")) | ||
if !errors.Is(err, errTestErrorAccumulatorWriteFailed) { | ||
t.Fatalf("Did not return error when write failed: %v", err) | ||
} | ||
} | ||
|
||
func TestErrorAccumulatorWriteErrors(t *testing.T) { | ||
var err error | ||
ts := test.NewTestServer().OpenAITestServer() | ||
ts.Start() | ||
defer ts.Close() | ||
|
||
config := DefaultConfig(test.GetTestToken()) | ||
config.BaseURL = ts.URL + "/v1" | ||
client := NewClientWithConfig(config) | ||
|
||
ctx := context.Background() | ||
|
||
stream, err := client.CreateChatCompletionStream(ctx, ChatCompletionRequest{}) | ||
if err != nil { | ||
t.Fatal(err) | ||
} | ||
stream.errAccumulator = &errorAccumulate{ | ||
buffer: &failingErrorBuffer{}, | ||
unmarshaler: &jsonUnmarshaler{}, | ||
} | ||
|
||
_, err = stream.Recv() | ||
if !errors.Is(err, errTestErrorAccumulatorWriteFailed) { | ||
t.Fatalf("Did not return error when write failed: %v", err) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.