From a0159ad2b00e4f127222814694bec68863395543 Mon Sep 17 00:00:00 2001
From: Mike Cutalo <mcutalo88@gmail.com>
Date: Tue, 7 Nov 2023 23:16:22 -0800
Subject: [PATCH] Support new fields for /v1/images/generation API (#530)

* add support for new image/generation api

* fix one lint

* add revised_prompt to response

* fix lints

* add CreateImageQualityStandard
---
 image.go          | 26 ++++++++++++++++++++++++--
 image_api_test.go |  9 ++++++++-
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/image.go b/image.go
index 4addcdb1e..4fe8b3a32 100644
--- a/image.go
+++ b/image.go
@@ -13,6 +13,9 @@ const (
 	CreateImageSize256x256   = "256x256"
 	CreateImageSize512x512   = "512x512"
 	CreateImageSize1024x1024 = "1024x1024"
+	// dall-e-3 supported only.
+	CreateImageSize1792x1024 = "1792x1024"
+	CreateImageSize1024x1792 = "1024x1792"
 )
 
 const (
@@ -20,11 +23,29 @@ const (
 	CreateImageResponseFormatB64JSON = "b64_json"
 )
 
+const (
+	CreateImageModelDallE2 = "dall-e-2"
+	CreateImageModelDallE3 = "dall-e-3"
+)
+
+const (
+	CreateImageQualityHD       = "hd"
+	CreateImageQualityStandard = "standard"
+)
+
+const (
+	CreateImageStyleVivid   = "vivid"
+	CreateImageStyleNatural = "natural"
+)
+
 // ImageRequest represents the request structure for the image API.
 type ImageRequest struct {
 	Prompt         string `json:"prompt,omitempty"`
+	Model          string `json:"model,omitempty"`
 	N              int    `json:"n,omitempty"`
+	Quality        string `json:"quality,omitempty"`
 	Size           string `json:"size,omitempty"`
+	Style          string `json:"style,omitempty"`
 	ResponseFormat string `json:"response_format,omitempty"`
 	User           string `json:"user,omitempty"`
 }
@@ -39,8 +60,9 @@ type ImageResponse struct {
 
 // ImageResponseDataInner represents a response data structure for image API.
 type ImageResponseDataInner struct {
-	URL     string `json:"url,omitempty"`
-	B64JSON string `json:"b64_json,omitempty"`
+	URL           string `json:"url,omitempty"`
+	B64JSON       string `json:"b64_json,omitempty"`
+	RevisedPrompt string `json:"revised_prompt,omitempty"`
 }
 
 // CreateImage - API call to create an image. This is the main endpoint of the DALL-E API.
diff --git a/image_api_test.go b/image_api_test.go
index 422f831fe..2eb46f2b4 100644
--- a/image_api_test.go
+++ b/image_api_test.go
@@ -19,7 +19,14 @@ func TestImages(t *testing.T) {
 	defer teardown()
 	server.RegisterHandler("/v1/images/generations", handleImageEndpoint)
 	_, err := client.CreateImage(context.Background(), openai.ImageRequest{
-		Prompt: "Lorem ipsum",
+		Prompt:         "Lorem ipsum",
+		Model:          openai.CreateImageModelDallE3,
+		N:              1,
+		Quality:        openai.CreateImageQualityHD,
+		Size:           openai.CreateImageSize1024x1024,
+		Style:          openai.CreateImageStyleVivid,
+		ResponseFormat: openai.CreateImageResponseFormatURL,
+		User:           "user",
 	})
 	checks.NoError(t, err, "CreateImage error")
 }