Add form builder (#235)

* 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

Author: sashabaranov

audio.go

@@ -4,8 +4,6 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"io"
-	"mime/multipart"
 	"net/http"
 	"os"
 )
@@ -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
 	}
 
@@ -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()
 }

audio_test.go

@@ -1,8 +1,9 @@
-package openai_test
+package openai //nolint:testpackage // testing private field
 
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"io"
 	"mime"
 	"mime/multipart"
@@ -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"
 
@@ -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")
+	}
+}

api.go renamed to client.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"strings"
 )
@@ -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.
@@ -26,6 +28,9 @@ func NewClientWithConfig(config ClientConfig) *Client {
 	return &Client{
 		config:         config,
 		requestBuilder: newRequestBuilder(),
+		createFormBuilder: func(body io.Writer) formBuilder {
+			return newFormBuilder(body)
+		},
 	}
 }
 

config.go

@@ -64,3 +64,7 @@ func DefaultAzureConfig(apiKey, baseURL, engine string) ClientConfig {
 		EmptyMessagesLimit: defaultEmptyMessagesLimit,
 	}
 }
+
+func (ClientConfig) String() string {
+	return "<OpenAI API ClientConfig>"
+}

files_test.go

@@ -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)

form_builder.go

@@ -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()
+}

image.go

@@ -3,8 +3,6 @@ package openai
 import (
 	"bytes"
 	"context"
-	"io"
-	"mime/multipart"
 	"net/http"
 	"os"
 	"strconv"
@@ -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
 }
@@ -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
 }