Added image methods

Author: djthorpe

pkg/openai/chat.go

@@ -12,29 +12,49 @@ const (
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
-func NewUserMessage(text string) Message {
-	return Message{
-		Role: "user", Content: &text,
+func NewMessage(role string, text ...string) *Message {
+	msg := &Message{
+		Role: role, Content: []MessageContent{},
 	}
+	return msg.AppendText(text...)
 }
 
-func NewSystemMessage(text string) Message {
-	return Message{
-		Role: "system", Content: &text,
+func (msg *Message) AppendText(text ...string) *Message {
+	for _, t := range text {
+		msg.Content = append(msg.Content, MessageContent{
+			Type: "text", Text: &t,
+		})
 	}
+	return msg
 }
 
-func NewAssistantMessage(text string) Message {
-	return Message{
-		Role: "assistant", Content: &text,
+func (msg *Message) AppendImageUrl(url ...string) *Message {
+	for _, v := range url {
+		msg.Content = append(msg.Content, MessageContent{
+			Type: "image_url", ImageUrl: &MessageContentImageUrl{
+				Url: v,
+			},
+		})
 	}
+	return msg
+}
+
+func (msg *Message) AppendImageFile(file ...string) *Message {
+	for _, v := range file {
+		msg.Content = append(msg.Content, MessageContent{
+			Type: "image_file", ImageFile: &MessageContentImageFile{
+				File: v,
+			},
+		})
+	}
+	return msg
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // API CALLS
 
 // Chat creates a model response for the given chat conversation.
-func (c *Client) Chat(messages []Message, opts ...Opt) (Chat, error) {
+func (c *Client) Chat(messages []*Message, opts ...Opt) (Chat, error) {
 	// Create the request
 	var request reqChat
 	request.Model = defaultChatCompletion

pkg/openai/chat.go_old

@@ -16,9 +16,8 @@ func Test_chat_001(t *testing.T) {
 	assert.NoError(err)
 	assert.NotNil(client)
 
-	response, err := client.Chat([]openai.Message{
-		openai.NewUserMessage("What would be the best app to use to get the weather in berlin today?"),
-	})
+	message := openai.NewMessage("user", "What would be the best app to use to get the weather in berlin today?")
+	response, err := client.Chat([]*openai.Message{message})
 	assert.NoError(err)
 	assert.NotNil(response)
 	assert.NotEmpty(response)
@@ -35,9 +34,8 @@ func Test_chat_002(t *testing.T) {
 	assert.NoError(err)
 	assert.NotNil(client)
 
-	response, err := client.Chat([]openai.Message{
-		openai.NewUserMessage("What is the weather in berlin today?"),
-	}, openai.OptFunction("get_weather", "Get the weather in a specific city and country", openai.ToolParameter{
+	message := openai.NewMessage("user", "What will the weather be like in Berlin tomorrow?")
+	response, err := client.Chat([]*openai.Message{message}, openai.OptFunction("get_weather", "Get the weather in a specific city and country", openai.ToolParameter{
 		Name:        "city",
 		Type:        "string",
 		Description: "The city to get the weather for",
@@ -69,10 +67,8 @@ func Test_chat_003(t *testing.T) {
 	assert.NoError(err)
 	assert.NotNil(client)
 
-	response, err := client.Chat([]openai.Message{
-		openai.NewUserMessage("What is in this image"),
-		openai.NewImageMessage("https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"),
-	}, openai.OptModel("gpt-4-vision-preview"))
+	message := openai.NewMessage("user", "What is in this image").AppendImageUrl("https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg")
+	response, err := client.Chat([]*openai.Message{message}, openai.OptModel("gpt-4-vision-preview"))
 	assert.NoError(err)
 	assert.NotNil(response)
 	assert.NotEmpty(response)

pkg/openai/chat_test.go

@@ -0,0 +1,41 @@
+package openai
+
+import (
+	// Packages
+	"github.com/mutablelogic/go-client/pkg/client"
+)
+
+///////////////////////////////////////////////////////////////////////////////
+// TYPES
+
+type responseImages struct {
+	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) {
+	var request reqImage
+	var response responseImages
+
+	// Create the request
+	request.Prompt = prompt
+	for _, opt := range opts {
+		if err := opt(&request); err != nil {
+			return nil, err
+		}
+	}
+
+	// Return the response
+	if payload, err := client.NewJSONRequest(request, client.ContentTypeJson); err != nil {
+		return nil, err
+	} else if err := c.Do(payload.Post(), &response, client.OptPath("images/generations")); err != nil {
+		return nil, err
+	}
+
+	// Return success
+	return response.Data, nil
+}

pkg/openai/image.go

@@ -0,0 +1,36 @@
+package openai_test
+
+import (
+	"os"
+	"testing"
+
+	opts "github.com/mutablelogic/go-client/pkg/client"
+	openai "github.com/mutablelogic/go-client/pkg/openai"
+	assert "github.com/stretchr/testify/assert"
+)
+
+func Test_image_001(t *testing.T) {
+	assert := assert.New(t)
+	client, err := openai.New(GetApiKey(t), opts.OptTrace(os.Stderr, true))
+	assert.NoError(err)
+	assert.NotNil(client)
+
+	images, err := client.CreateImages("A painting of a cat", openai.OptCount(1))
+	assert.NoError(err)
+	assert.NotNil(images)
+	assert.NotEmpty(images)
+	assert.Len(images, 1)
+}
+
+func Test_image_002(t *testing.T) {
+	assert := assert.New(t)
+	client, err := openai.New(GetApiKey(t), opts.OptTrace(os.Stderr, true))
+	assert.NoError(err)
+	assert.NotNil(client)
+
+	images, err := client.CreateImages("A painting of a cat", openai.OptResponseFormat("b64_json"), openai.OptCount(1))
+	assert.NoError(err)
+	assert.NotNil(images)
+	assert.NotEmpty(images)
+	assert.Len(images, 1)
+}