Skip to content

Commit

Permalink
feat: support direct bytes for file upload (sashabaranov#568)
Browse files Browse the repository at this point in the history
* feat: support direct bytes for file upload

* add test for errors

* add coverage
  • Loading branch information
floodfx authored Nov 15, 2023
1 parent fe67abb commit 71848cc
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 0 deletions.
3 changes: 3 additions & 0 deletions client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,9 @@ func TestClientReturnsRequestBuilderErrors(t *testing.T) {
{"CreateImage", func() (any, error) {
return client.CreateImage(ctx, ImageRequest{})
}},
{"CreateFileBytes", func() (any, error) {
return client.CreateFileBytes(ctx, FileBytesRequest{})
}},
{"DeleteFile", func() (any, error) {
return nil, client.DeleteFile(ctx, "")
}},
Expand Down
49 changes: 49 additions & 0 deletions files.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,24 @@ type FileRequest struct {
Purpose string `json:"purpose"`
}

// PurposeType represents the purpose of the file when uploading.
type PurposeType string

const (
PurposeFineTune PurposeType = "fine-tune"
PurposeAssistants PurposeType = "assistants"
)

// FileBytesRequest represents a file upload request.
type FileBytesRequest struct {
// the name of the uploaded file in OpenAI
Name string
// the bytes of the file
Bytes []byte
// the purpose of the file
Purpose PurposeType
}

// File struct represents an OpenAPI file.
type File struct {
Bytes int `json:"bytes"`
Expand All @@ -36,6 +54,37 @@ type FilesList struct {
httpHeader
}

// CreateFileBytes uploads bytes directly to OpenAI without requiring a local file.
func (c *Client) CreateFileBytes(ctx context.Context, request FileBytesRequest) (file File, err error) {
var b bytes.Buffer
reader := bytes.NewReader(request.Bytes)
builder := c.createFormBuilder(&b)

err = builder.WriteField("purpose", string(request.Purpose))
if err != nil {
return
}

err = builder.CreateFormFileReader("file", reader, request.Name)
if err != nil {
return
}

err = builder.Close()
if err != nil {
return
}

req, err := c.newRequest(ctx, http.MethodPost, c.fullURL("/files"),
withBody(&b), withContentType(builder.FormDataContentType()))
if err != nil {
return
}

err = c.sendRequest(req, &file)
return
}

// CreateFile uploads a jsonl file to GPT3
// FilePath must be a local file path.
func (c *Client) CreateFile(ctx context.Context, request FileRequest) (file File, err error) {
Expand Down
13 changes: 13 additions & 0 deletions files_api_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,19 @@ import (
"github.com/sashabaranov/go-openai/internal/test/checks"
)

func TestFileBytesUpload(t *testing.T) {
client, server, teardown := setupOpenAITestServer()
defer teardown()
server.RegisterHandler("/v1/files", handleCreateFile)
req := openai.FileBytesRequest{
Name: "foo",
Bytes: []byte("foo"),
Purpose: openai.PurposeFineTune,
}
_, err := client.CreateFileBytes(context.Background(), req)
checks.NoError(t, err, "CreateFile error")
}

func TestFileUpload(t *testing.T) {
client, server, teardown := setupOpenAITestServer()
defer teardown()
Expand Down
50 changes: 50 additions & 0 deletions files_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,53 @@ import (
"github.com/sashabaranov/go-openai/internal/test/checks"
)

func TestFileBytesUploadWithFailingFormBuilder(t *testing.T) {
config := DefaultConfig("")
config.BaseURL = ""
client := NewClientWithConfig(config)
mockBuilder := &mockFormBuilder{}
client.createFormBuilder = func(io.Writer) utils.FormBuilder {
return mockBuilder
}

ctx := context.Background()
req := FileBytesRequest{
Name: "foo",
Bytes: []byte("foo"),
Purpose: PurposeAssistants,
}

mockError := fmt.Errorf("mockWriteField error")
mockBuilder.mockWriteField = func(string, string) error {
return mockError
}
_, err := client.CreateFileBytes(ctx, req)
checks.ErrorIs(t, err, mockError, "CreateFile should return error if form builder fails")

mockError = fmt.Errorf("mockCreateFormFile error")
mockBuilder.mockWriteField = func(string, string) error {
return nil
}
mockBuilder.mockCreateFormFileReader = func(string, io.Reader, string) error {
return mockError
}
_, err = client.CreateFileBytes(ctx, req)
checks.ErrorIs(t, err, mockError, "CreateFile should return error if form builder fails")

mockError = fmt.Errorf("mockClose error")
mockBuilder.mockWriteField = func(string, string) error {
return nil
}
mockBuilder.mockCreateFormFileReader = func(string, io.Reader, string) error {
return nil
}
mockBuilder.mockClose = func() error {
return mockError
}
_, err = client.CreateFileBytes(ctx, req)
checks.ErrorIs(t, err, mockError, "CreateFile should return error if form builder fails")
}

func TestFileUploadWithFailingFormBuilder(t *testing.T) {
config := DefaultConfig("")
config.BaseURL = ""
Expand Down Expand Up @@ -55,6 +102,9 @@ func TestFileUploadWithFailingFormBuilder(t *testing.T) {
return mockError
}
_, err = client.CreateFile(ctx, req)
if err == nil {
t.Fatal("CreateFile should return error if form builder fails")
}
checks.ErrorIs(t, err, mockError, "CreateFile should return error if form builder fails")
}

Expand Down

0 comments on commit 71848cc

Please sign in to comment.