move error_accumulator into internal pkg (#304) (#335)

* move error_accumulator into internal pkg (#304)

* move error_accumulator into internal pkg (#304)

* add a test for ErrTooManyEmptyStreamMessages in stream_reader (#304)

Author: vvatanabe

chat_stream.go

@@ -66,7 +66,7 @@ func (c *Client) CreateChatCompletionStream(
 			emptyMessagesLimit: c.config.EmptyMessagesLimit,
 			reader:             bufio.NewReader(resp.Body),
 			response:           resp,
-			errAccumulator:     newErrorAccumulator(),
+			errAccumulator:     utils.NewErrorAccumulator(),
 			unmarshaler:        &utils.JSONUnmarshaler{},
 		},
 	}

chat_stream_test.go

@@ -1,7 +1,7 @@
-package openai_test
+package openai //nolint:testpackage // testing private field
 
 import (
-	. "github.com/sashabaranov/go-openai"
+	utils "github.com/sashabaranov/go-openai/internal"
 	"github.com/sashabaranov/go-openai/internal/test"
 	"github.com/sashabaranov/go-openai/internal/test/checks"
 
@@ -63,9 +63,9 @@ func TestCreateChatCompletionStream(t *testing.T) {
 	// Client portion of the test
 	config := DefaultConfig(test.GetTestToken())
 	config.BaseURL = server.URL + "/v1"
-	config.HTTPClient.Transport = &tokenRoundTripper{
-		test.GetTestToken(),
-		http.DefaultTransport,
+	config.HTTPClient.Transport = &test.TokenRoundTripper{
+		Token:    test.GetTestToken(),
+		Fallback: http.DefaultTransport,
 	}
 
 	client := NewClientWithConfig(config)
@@ -170,9 +170,9 @@ func TestCreateChatCompletionStreamError(t *testing.T) {
 	// Client portion of the test
 	config := DefaultConfig(test.GetTestToken())
 	config.BaseURL = server.URL + "/v1"
-	config.HTTPClient.Transport = &tokenRoundTripper{
-		test.GetTestToken(),
-		http.DefaultTransport,
+	config.HTTPClient.Transport = &test.TokenRoundTripper{
+		Token:    test.GetTestToken(),
+		Fallback: http.DefaultTransport,
 	}
 
 	client := NewClientWithConfig(config)
@@ -227,9 +227,9 @@ func TestCreateChatCompletionStreamRateLimitError(t *testing.T) {
 	// Client portion of the test
 	config := DefaultConfig(test.GetTestToken())
 	config.BaseURL = ts.URL + "/v1"
-	config.HTTPClient.Transport = &tokenRoundTripper{
-		test.GetTestToken(),
-		http.DefaultTransport,
+	config.HTTPClient.Transport = &test.TokenRoundTripper{
+		Token:    test.GetTestToken(),
+		Fallback: http.DefaultTransport,
 	}
 
 	client := NewClientWithConfig(config)
@@ -255,6 +255,33 @@ func TestCreateChatCompletionStreamRateLimitError(t *testing.T) {
 	t.Logf("%+v\n", apiErr)
 }
 
+func TestCreateChatCompletionStreamErrorAccumulatorWriteErrors(t *testing.T) {
+	var err error
+	server := test.NewTestServer()
+	server.RegisterHandler("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "error", 200)
+	})
+	ts := server.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{})
+	checks.NoError(t, err)
+
+	stream.errAccumulator = &utils.DefaultErrorAccumulator{
+		Buffer: &test.FailingErrorBuffer{},
+	}
+
+	_, err = stream.Recv()
+	checks.ErrorIs(t, err, test.ErrTestErrorAccumulatorWriteFailed, "Did not return error when Write failed", err.Error())
+}
+
 // Helper funcs.
 func compareChatResponses(r1, r2 ChatCompletionStreamResponse) bool {
 	if r1.ID != r2.ID || r1.Object != r2.Object || r1.Created != r2.Created || r1.Model != r2.Model {

error_accumulator.go

@@ -0,0 +1,44 @@
+package openai
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+)
+
+type ErrorAccumulator interface {
+	Write(p []byte) error
+	Bytes() []byte
+}
+
+type errorBuffer interface {
+	io.Writer
+	Len() int
+	Bytes() []byte
+}
+
+type DefaultErrorAccumulator struct {
+	Buffer errorBuffer
+}
+
+func NewErrorAccumulator() ErrorAccumulator {
+	return &DefaultErrorAccumulator{
+		Buffer: &bytes.Buffer{},
+	}
+}
+
+func (e *DefaultErrorAccumulator) 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 *DefaultErrorAccumulator) Bytes() (errBytes []byte) {
+	if e.Buffer.Len() == 0 {
+		return
+	}
+	errBytes = e.Buffer.Bytes()
+	return
+}

error_accumulator_test.go

@@ -0,0 +1,41 @@
+package openai_test
+
+import (
+	"bytes"
+	"errors"
+	"testing"
+
+	utils "github.com/sashabaranov/go-openai/internal"
+	"github.com/sashabaranov/go-openai/internal/test"
+)
+
+func TestErrorAccumulatorBytes(t *testing.T) {
+	accumulator := &utils.DefaultErrorAccumulator{
+		Buffer: &bytes.Buffer{},
+	}
+
+	errBytes := accumulator.Bytes()
+	if len(errBytes) != 0 {
+		t.Fatalf("Did not return nil with empty bytes: %s", string(errBytes))
+	}
+
+	err := accumulator.Write([]byte("{}"))
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+
+	errBytes = accumulator.Bytes()
+	if len(errBytes) == 0 {
+		t.Fatalf("Did not return error bytes when has error: %s", string(errBytes))
+	}
+}
+
+func TestErrorByteWriteErrors(t *testing.T) {
+	accumulator := &utils.DefaultErrorAccumulator{
+		Buffer: &test.FailingErrorBuffer{},
+	}
+	err := accumulator.Write([]byte("{"))
+	if !errors.Is(err, test.ErrTestErrorAccumulatorWriteFailed) {
+		t.Fatalf("Did not return error when write failed: %v", err)
+	}
+}

internal/error_accumulator.go

@@ -0,0 +1,21 @@
+package test
+
+import "errors"
+
+var (
+	ErrTestErrorAccumulatorWriteFailed = errors.New("test error accumulator failed")
+)
+
+type 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{}
+}

internal/error_accumulator_test.go

@@ -3,6 +3,7 @@ package test
 import (
 	"github.com/sashabaranov/go-openai/internal/test/checks"
 
+	"net/http"
 	"os"
 	"testing"
 )
@@ -27,3 +28,26 @@ func CreateTestDirectory(t *testing.T) (path string, cleanup func()) {
 
 	return path, func() { os.RemoveAll(path) }
 }
+
+// TokenRoundTripper is a struct that implements the RoundTripper
+// interface, specifically to handle the authentication token by adding a token
+// to the request header. We need this because the API requires that each
+// request include a valid API token in the headers for authentication and
+// authorization.
+type TokenRoundTripper struct {
+	Token    string
+	Fallback http.RoundTripper
+}
+
+// RoundTrip takes an *http.Request as input and returns an
+// *http.Response and an error.
+//
+// It is expected to use the provided request to create a connection to an HTTP
+// server and return the response, or an error if one occurred. The returned
+// Response should have its Body closed. If the RoundTrip method returns an
+// error, the Client's Get, Head, Post, and PostForm methods return the same
+// error.
+func (t *TokenRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	req.Header.Set("Authorization", "Bearer "+t.Token)
+	return t.Fallback.RoundTrip(req)
+}

internal/test/failer.go

@@ -55,7 +55,7 @@ func (c *Client) CreateCompletionStream(
 			emptyMessagesLimit: c.config.EmptyMessagesLimit,
 			reader:             bufio.NewReader(resp.Body),
 			response:           resp,
-			errAccumulator:     newErrorAccumulator(),
+			errAccumulator:     utils.NewErrorAccumulator(),
 			unmarshaler:        &utils.JSONUnmarshaler{},
 		},
 	}