Update images

Author: djthorpe

pkg/openai/image.go

@@ -1,23 +1,29 @@
 package openai
 
 import (
+	"encoding/base64"
+	"io"
+
 	// Packages
 	"github.com/mutablelogic/go-client/pkg/client"
+
+	// Namespace imports
+	. "github.com/djthorpe/go-errors"
 )
 
 ///////////////////////////////////////////////////////////////////////////////
 // TYPES
 
 type responseImages struct {
-	Created int64   `json:"created"`
-	Data    []Image `json:"data"`
+	Created int64    `json:"created"`
+	Data    []*Image `json:"data"`
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // API CALLS
 
 // CreateImage generates one or more images from a prompt
-func (c *Client) CreateImages(prompt string, opts ...Opt) ([]Image, error) {
+func (c *Client) CreateImages(prompt string, opts ...Opt) ([]*Image, error) {
 	var request reqImage
 	var response responseImages
 
@@ -39,3 +45,23 @@ func (c *Client) CreateImages(prompt string, opts ...Opt) ([]Image, error) {
 	// Return success
 	return response.Data, nil
 }
+
+// WriteImage writes an image and returns the number of bytes written
+func (c *Client) WriteImage(w io.Writer, image *Image) (int, error) {
+	if image == nil {
+		return 0, ErrBadParameter.With("WriteImage")
+	}
+	// Handle url or data
+	switch {
+	case image.Data != "":
+		if data, err := base64.StdEncoding.DecodeString(image.Data); err != nil {
+			return 0, err
+		} else if n, err := w.Write(data); err != nil {
+			return 0, err
+		} else {
+			return n, nil
+		}
+	default:
+		return 0, ErrNotImplemented.With("WriteImage")
+	}
+}

pkg/openai/image_test.go

@@ -1,7 +1,9 @@
 package openai_test
 
 import (
+	"fmt"
 	"os"
+	"path/filepath"
 	"testing"
 
 	opts "github.com/mutablelogic/go-client/pkg/client"
@@ -28,9 +30,25 @@ func Test_image_002(t *testing.T) {
 	assert.NoError(err)
 	assert.NotNil(client)
 
-	images, err := client.CreateImages("A painting of a cat", openai.OptResponseFormat("b64_json"), openai.OptCount(1))
+	// Create one image
+	images, err := client.CreateImages("A painting of a cat in the style of Salvador Dali", openai.OptResponseFormat("b64_json"), openai.OptCount(1))
 	assert.NoError(err)
 	assert.NotNil(images)
 	assert.NotEmpty(images)
 	assert.Len(images, 1)
+
+	// Output images
+	for n, image := range images {
+		filename := filepath.Join(os.TempDir(), fmt.Sprintf("test-%d.png", n))
+		if w, err := os.Create(filename); err != nil {
+			t.Error(err)
+		} else {
+			defer w.Close()
+			t.Log("Writing", w.Name())
+			n, err := client.WriteImage(w, image)
+			assert.NoError(err)
+			assert.NotZero(n)
+		}
+	}
+
 }