From 59603f4b43e61f7d2be3d12cff3ea8fdd524fd7d Mon Sep 17 00:00:00 2001
From: csuzhang <csuzhang@tencent.com>
Date: Tue, 28 Mar 2023 14:36:03 +0800
Subject: [PATCH] feat: support image generation

---
 README-CN.md |  1 +
 README.md    |  1 +
 gpt.go       | 17 +++++++++++++++++
 model.go     | 21 +++++++++++++++++++++
 4 files changed, 40 insertions(+)

diff --git a/README-CN.md b/README-CN.md
index dc3947f..05d0b5b 100644
--- a/README-CN.md
+++ b/README-CN.md
@@ -48,6 +48,7 @@ make chatgpt-example
 - [x] Completion API (是主要的 gpt-3 API)
 - [x] 对 Completion API 的流式支持
 - [x] 文档搜索 API
+- [x] 图片生成 API
 - [x] 替换默认 url、用户代理、超时和其他选项
 
 ## 贡献者
diff --git a/README.md b/README.md
index 9a04456..d5c44ee 100644
--- a/README.md
+++ b/README.md
@@ -52,6 +52,7 @@ Check out the go docs for more detailed documentation on the types and methods p
 - [x] Completion API (this is the main gpt-3 API)
 - [x] Streaming support for the Completion API
 - [x] Document Search API
+- [x] Image generation API
 - [x] Overriding default url, user-agent, timeout, and other options
 
 ## Contributor
diff --git a/gpt.go b/gpt.go
index bcde523..cc843a9 100644
--- a/gpt.go
+++ b/gpt.go
@@ -362,6 +362,23 @@ func (c *client) Embeddings(ctx context.Context, request EmbeddingsRequest) (*Em
 	return &output, nil
 }
 
+// Image creates an image
+func (c *client) Image(ctx context.Context, request ImageRequest) (*ImageResponse, error) {
+	req, err := c.newRequest(ctx, "POST", "/images/generations", request)
+	if err != nil {
+		return nil, err
+	}
+	rsp, err := c.performRequest(req)
+	if err != nil {
+		return nil, err
+	}
+	output := ImageResponse{}
+	if err := getResponseObject(rsp, &output); err != nil {
+		return nil, err
+	}
+	return &output, nil
+}
+
 func (c *client) performRequest(req *http.Request) (*http.Response, error) {
 	rsp, err := c.httpClient.Do(req)
 	if err != nil {
diff --git a/model.go b/model.go
index 47c8b46..687e377 100644
--- a/model.go
+++ b/model.go
@@ -285,3 +285,24 @@ type SearchResponse struct {
 	Data   []SearchData `json:"data"`
 	Object string       `json:"object"`
 }
+
+// ImageRequest represents the request structure for the image API.
+type ImageRequest struct {
+	Prompt         string `json:"prompt,omitempty"`
+	N              int    `json:"n,omitempty"`
+	Size           string `json:"size,omitempty"`
+	ResponseFormat string `json:"response_format,omitempty"`
+	User           string `json:"user,omitempty"`
+}
+
+// ImageResponse represents a response structure for image API.
+type ImageResponse struct {
+	Created int64                    `json:"created,omitempty"`
+	Data    []ImageResponseDataInner `json:"data,omitempty"`
+}
+
+// ImageResponseDataInner represents a response data structure for image API.
+type ImageResponseDataInner struct {
+	URL     string `json:"url,omitempty"`
+	B64JSON string `json:"b64_json,omitempty"`
+}