Skip to content

Commit

Permalink
Add form builder (#235)
Browse files Browse the repository at this point in the history
* add form builder

* cover VariImage

* test for closing errors

* simplify tests

* add audio api test coverage

* don't leak authToken when printed

* rename api->client

* fix test
  • Loading branch information
sashabaranov authored Apr 8, 2023
1 parent 2f3700f commit 226ff32
Show file tree
Hide file tree
Showing 8 changed files with 270 additions and 70 deletions.
48 changes: 10 additions & 38 deletions audio.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,6 @@ import (
"bytes"
"context"
"fmt"
"io"
"mime/multipart"
"net/http"
"os"
)
Expand Down Expand Up @@ -55,9 +53,9 @@ func (c *Client) callAudioAPI(
endpointSuffix string,
) (response AudioResponse, err error) {
var formBody bytes.Buffer
w := multipart.NewWriter(&formBody)
builder := c.createFormBuilder(&formBody)

if err = audioMultipartForm(request, w); err != nil {
if err = audioMultipartForm(request, builder); err != nil {
return
}

Expand All @@ -66,81 +64,55 @@ func (c *Client) callAudioAPI(
if err != nil {
return
}
req.Header.Add("Content-Type", w.FormDataContentType())
req.Header.Add("Content-Type", builder.formDataContentType())

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

// audioMultipartForm creates a form with audio file contents and the name of the model to use for
// audio processing.
func audioMultipartForm(request AudioRequest, w *multipart.Writer) error {
func audioMultipartForm(request AudioRequest, b formBuilder) error {
f, err := os.Open(request.FilePath)
if err != nil {
return fmt.Errorf("opening audio file: %w", err)
}
defer f.Close()

fw, err := w.CreateFormFile("file", f.Name())
err = b.createFormFile("file", f)
if err != nil {
return fmt.Errorf("creating form file: %w", err)
}

if _, err = io.Copy(fw, f); err != nil {
return fmt.Errorf("reading from opened audio file: %w", err)
}

fw, err = w.CreateFormField("model")
err = b.writeField("model", request.Model)
if err != nil {
return fmt.Errorf("creating form field: %w", err)
}

modelName := bytes.NewReader([]byte(request.Model))
if _, err = io.Copy(fw, modelName); err != nil {
return fmt.Errorf("writing model name: %w", err)
}

// Create a form field for the prompt (if provided)
if request.Prompt != "" {
fw, err = w.CreateFormField("prompt")
err = b.writeField("prompt", request.Prompt)
if err != nil {
return fmt.Errorf("creating form field: %w", err)
}

prompt := bytes.NewReader([]byte(request.Prompt))
if _, err = io.Copy(fw, prompt); err != nil {
return fmt.Errorf("writing prompt: %w", err)
}
}

// Create a form field for the temperature (if provided)
if request.Temperature != 0 {
fw, err = w.CreateFormField("temperature")
err = b.writeField("temperature", fmt.Sprintf("%.2f", request.Temperature))
if err != nil {
return fmt.Errorf("creating form field: %w", err)
}

temperature := bytes.NewReader([]byte(fmt.Sprintf("%.2f", request.Temperature)))
if _, err = io.Copy(fw, temperature); err != nil {
return fmt.Errorf("writing temperature: %w", err)
}
}

// Create a form field for the language (if provided)
if request.Language != "" {
fw, err = w.CreateFormField("language")
err = b.writeField("language", request.Language)
if err != nil {
return fmt.Errorf("creating form field: %w", err)
}

language := bytes.NewReader([]byte(request.Language))
if _, err = io.Copy(fw, language); err != nil {
return fmt.Errorf("writing language: %w", err)
}
}

// Close the multipart writer
w.Close()

return nil
return b.close()
}
48 changes: 46 additions & 2 deletions audio_test.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
package openai_test
package openai //nolint:testpackage // testing private field

import (
"bytes"
"errors"
"fmt"
"io"
"mime"
"mime/multipart"
Expand All @@ -11,7 +12,6 @@ import (
"path/filepath"
"strings"

. "github.com/sashabaranov/go-openai"
"github.com/sashabaranov/go-openai/internal/test"
"github.com/sashabaranov/go-openai/internal/test/checks"

Expand Down Expand Up @@ -188,3 +188,47 @@ func handleAudioEndpoint(w http.ResponseWriter, r *http.Request) {
return
}
}

func TestAudioWithFailingFormBuilder(t *testing.T) {
dir, cleanup := createTestDirectory(t)
defer cleanup()
path := filepath.Join(dir, "fake.mp3")
createTestFile(t, path)

req := AudioRequest{
FilePath: path,
Prompt: "test",
Temperature: 0.5,
Language: "en",
}

mockFailedErr := fmt.Errorf("mock form builder fail")
mockBuilder := &mockFormBuilder{}

mockBuilder.mockCreateFormFile = func(string, *os.File) error {
return mockFailedErr
}
err := audioMultipartForm(req, mockBuilder)
checks.ErrorIs(t, err, mockFailedErr, "audioMultipartForm should return error if form builder fails")

mockBuilder.mockCreateFormFile = func(string, *os.File) error {
return nil
}

var failForField string
mockBuilder.mockWriteField = func(fieldname, value string) error {
if fieldname == failForField {
return mockFailedErr
}
return nil
}

failOn := []string{"model", "prompt", "temperature", "language"}
for _, failingField := range failOn {
failForField = failingField
mockFailedErr = fmt.Errorf("mock form builder fail on field %s", failingField)

err = audioMultipartForm(req, mockBuilder)
checks.ErrorIs(t, err, mockFailedErr, "audioMultipartForm should return error if form builder fails")
}
}
7 changes: 6 additions & 1 deletion api.go → client.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
)
Expand All @@ -12,7 +13,8 @@ import (
type Client struct {
config ClientConfig

requestBuilder requestBuilder
requestBuilder requestBuilder
createFormBuilder func(io.Writer) formBuilder
}

// NewClient creates new OpenAI API client.
Expand All @@ -26,6 +28,9 @@ func NewClientWithConfig(config ClientConfig) *Client {
return &Client{
config: config,
requestBuilder: newRequestBuilder(),
createFormBuilder: func(body io.Writer) formBuilder {
return newFormBuilder(body)
},
}
}

Expand Down
4 changes: 4 additions & 0 deletions config.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,3 +64,7 @@ func DefaultAzureConfig(apiKey, baseURL, engine string) ClientConfig {
EmptyMessagesLimit: defaultEmptyMessagesLimit,
}
}

func (ClientConfig) String() string {
return "<OpenAI API ClientConfig>"
}
2 changes: 1 addition & 1 deletion files_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ func TestFileUpload(t *testing.T) {

req := FileRequest{
FileName: "test.go",
FilePath: "api.go",
FilePath: "client.go",
Purpose: "fine-tune",
}
_, err = client.CreateFile(ctx, req)
Expand Down
49 changes: 49 additions & 0 deletions form_builder.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
package openai

import (
"io"
"mime/multipart"
"os"
)

type formBuilder interface {
createFormFile(fieldname string, file *os.File) error
writeField(fieldname, value string) error
close() error
formDataContentType() string
}

type defaultFormBuilder struct {
writer *multipart.Writer
}

func newFormBuilder(body io.Writer) *defaultFormBuilder {
return &defaultFormBuilder{
writer: multipart.NewWriter(body),
}
}

func (fb *defaultFormBuilder) createFormFile(fieldname string, file *os.File) error {
fieldWriter, err := fb.writer.CreateFormFile(fieldname, file.Name())
if err != nil {
return err
}

_, err = io.Copy(fieldWriter, file)
if err != nil {
return err
}
return nil
}

func (fb *defaultFormBuilder) writeField(fieldname, value string) error {
return fb.writer.WriteField(fieldname, value)
}

func (fb *defaultFormBuilder) close() error {
return fb.writer.Close()
}

func (fb *defaultFormBuilder) formDataContentType() string {
return fb.writer.FormDataContentType()
}
46 changes: 20 additions & 26 deletions image.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,6 @@ package openai
import (
"bytes"
"context"
"io"
"mime/multipart"
"net/http"
"os"
"strconv"
Expand Down Expand Up @@ -67,50 +65,46 @@ type ImageEditRequest struct {
// CreateEditImage - API call to create an image. This is the main endpoint of the DALL-E API.
func (c *Client) CreateEditImage(ctx context.Context, request ImageEditRequest) (response ImageResponse, err error) {
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
builder := c.createFormBuilder(body)

// image
image, err := writer.CreateFormFile("image", request.Image.Name())
if err != nil {
return
}
_, err = io.Copy(image, request.Image)
err = builder.createFormFile("image", request.Image)
if err != nil {
return
}

// mask, it is optional
if request.Mask != nil {
mask, err2 := writer.CreateFormFile("mask", request.Mask.Name())
if err2 != nil {
return
}
_, err = io.Copy(mask, request.Mask)
err = builder.createFormFile("mask", request.Mask)
if err != nil {
return
}
}

err = writer.WriteField("prompt", request.Prompt)
err = builder.writeField("prompt", request.Prompt)
if err != nil {
return
}
err = builder.writeField("n", strconv.Itoa(request.N))
if err != nil {
return
}
err = writer.WriteField("n", strconv.Itoa(request.N))
err = builder.writeField("size", request.Size)
if err != nil {
return
}
err = writer.WriteField("size", request.Size)
err = builder.close()
if err != nil {
return
}
writer.Close()

urlSuffix := "/images/edits"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.fullURL(urlSuffix), body)
if err != nil {
return
}

req.Header.Set("Content-Type", writer.FormDataContentType())
req.Header.Set("Content-Type", builder.formDataContentType())
err = c.sendRequest(req, &response)
return
}
Expand All @@ -126,35 +120,35 @@ type ImageVariRequest struct {
// Use abbreviations(vari for variation) because ci-lint has a single-line length limit ...
func (c *Client) CreateVariImage(ctx context.Context, request ImageVariRequest) (response ImageResponse, err error) {
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
builder := c.createFormBuilder(body)

// image
image, err := writer.CreateFormFile("image", request.Image.Name())
err = builder.createFormFile("image", request.Image)
if err != nil {
return
}
_, err = io.Copy(image, request.Image)

err = builder.writeField("n", strconv.Itoa(request.N))
if err != nil {
return
}

err = writer.WriteField("n", strconv.Itoa(request.N))
err = builder.writeField("size", request.Size)
if err != nil {
return
}
err = writer.WriteField("size", request.Size)
err = builder.close()
if err != nil {
return
}
writer.Close()

//https://platform.openai.com/docs/api-reference/images/create-variation
urlSuffix := "/images/variations"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.fullURL(urlSuffix), body)
if err != nil {
return
}

req.Header.Set("Content-Type", writer.FormDataContentType())
req.Header.Set("Content-Type", builder.formDataContentType())
err = c.sendRequest(req, &response)
return
}
Loading

0 comments on commit 226ff32

Please sign in to comment.