feat: add support for image editing

Author: kardolus

api/client/client.go

@@ -17,6 +17,7 @@ import (
 	"golang.org/x/text/language"
 	"io"
 	"mime/multipart"
+	"net/textproto"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -399,22 +400,115 @@ func (c *Client) SynthesizeSpeech(inputText, outputPath string) error {
 // Returns:
 //   - An error if any part of the request, decoding, or file writing fails.
 func (c *Client) GenerateImage(inputText, outputPath string) error {
-	reqBody := api.Draw{
+	req := api.Draw{
 		Model:  c.Config.Model,
 		Prompt: inputText,
 	}
 
-	body, err := json.Marshal(reqBody)
+	return c.postAndWriteBinaryOutput(
+		c.getEndpoint(c.Config.ImageGenerationsPath),
+		req,
+		outputPath,
+		"image",
+		func(respBytes []byte) ([]byte, error) {
+			var response struct {
+				Data []struct {
+					B64 string `json:"b64_json"`
+				} `json:"data"`
+			}
+			if err := json.Unmarshal(respBytes, &response); err != nil {
+				return nil, fmt.Errorf("failed to decode response: %w", err)
+			}
+			if len(response.Data) == 0 {
+				return nil, fmt.Errorf("no image data returned")
+			}
+			decoded, err := base64.StdEncoding.DecodeString(response.Data[0].B64)
+			if err != nil {
+				return nil, fmt.Errorf("failed to decode base64 image: %w", err)
+			}
+			return decoded, nil
+		},
+	)
+}
+
+// EditImage edits an input image using a text prompt and writes the modified image to the specified output path.
+//
+// This method sends a multipart/form-data POST request to the image editing endpoint
+// (typically OpenAI's /v1/images/edits). The request includes:
+//   - The image file to edit.
+//   - A text prompt describing how the image should be modified.
+//   - The model ID (e.g., gpt-image-1).
+//
+// The response is expected to contain a base64-encoded image, which is decoded and written to the outputPath.
+//
+// Parameters:
+//   - inputText: A text prompt describing the desired modifications to the image.
+//   - inputPath: The file path to the source image (must be a supported format: PNG, JPEG, or WebP).
+//   - outputPath: The file path where the edited image will be saved.
+//
+// Returns:
+//   - An error if any step of the process fails: reading the file, building the request, sending it,
+//     decoding the response, or writing the output image.
+//
+// Example:
+//
+//	err := client.EditImage("Add a rainbow in the sky", "input.png", "output.png")
+//	if err != nil {
+//	    log.Fatal(err)
+//	}
+func (c *Client) EditImage(inputText, inputPath, outputPath string) error {
+	endpoint := c.getEndpoint(c.Config.ImageEditsPath)
+
+	file, err := c.reader.Open(inputPath)
 	if err != nil {
-		return fmt.Errorf("failed to marshal request: %w", err)
+		return fmt.Errorf("failed to open input image: %w", err)
 	}
+	defer file.Close()
 
-	endpoint := c.getEndpoint(c.Config.DrawPath)
-	c.printRequestDebugInfo(endpoint, body, nil)
+	var buf bytes.Buffer
+	writer := multipart.NewWriter(&buf)
 
-	respBytes, err := c.caller.Post(endpoint, body, false)
+	mimeType, err := c.getMimeTypeFromFileContent(inputPath)
+	if err != nil {
+		return fmt.Errorf("failed to detect MIME type: %w", err)
+	}
+	if !strings.HasPrefix(mimeType, "image/") {
+		return fmt.Errorf("unsupported MIME type: %s", mimeType)
+	}
+
+	header := make(textproto.MIMEHeader)
+	header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="image"; filename="%s"`, filepath.Base(inputPath)))
+	header.Set("Content-Type", mimeType)
+
+	part, err := writer.CreatePart(header)
+	if err != nil {
+		return fmt.Errorf("failed to create image part: %w", err)
+	}
+	if _, err := io.Copy(part, file); err != nil {
+		return fmt.Errorf("failed to copy image data: %w", err)
+	}
+
+	if err := writer.WriteField("prompt", inputText); err != nil {
+		return fmt.Errorf("failed to add prompt: %w", err)
+	}
+	if err := writer.WriteField("model", c.Config.Model); err != nil {
+		return fmt.Errorf("failed to add model: %w", err)
+	}
+
+	if err := writer.Close(); err != nil {
+		return fmt.Errorf("failed to close multipart writer: %w", err)
+	}
+
+	c.printRequestDebugInfo(endpoint, buf.Bytes(), map[string]string{
+		"Content-Type": writer.FormDataContentType(),
+	})
+
+	respBytes, err := c.caller.PostWithHeaders(endpoint, buf.Bytes(), map[string]string{
+		c.Config.AuthHeader: fmt.Sprintf("%s %s", c.Config.AuthTokenPrefix, c.Config.APIKey),
+		"Content-Type":      writer.FormDataContentType(),
+	})
 	if err != nil {
-		return fmt.Errorf("failed to generate image: %w", err)
+		return fmt.Errorf("failed to edit image: %w", err)
 	}
 
 	// Parse the JSON and extract b64_json

api/client/client_test.go

@@ -1,6 +1,7 @@
 package client_test
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -16,6 +17,7 @@ import (
 	"github.com/kardolus/chatgpt-cli/history"
 	"github.com/kardolus/chatgpt-cli/internal"
 	"github.com/kardolus/chatgpt-cli/test"
+	"io"
 	"os"
 	"strings"
 	"testing"
@@ -1060,7 +1062,7 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 		})
 		it("throws an error when the http call fails", func() {
 			mockCaller.EXPECT().
-				Post(subject.Config.URL+subject.Config.DrawPath, body, false).
+				Post(subject.Config.URL+subject.Config.ImageGenerationsPath, body, false).
 				Return(nil, errors.New(errorText))
 
 			err := subject.GenerateImage(inputText, outputFile)
@@ -1069,7 +1071,7 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 		})
 		it("throws an error when no image data is returned", func() {
 			mockCaller.EXPECT().
-				Post(subject.Config.URL+subject.Config.DrawPath, body, false).
+				Post(subject.Config.URL+subject.Config.ImageGenerationsPath, body, false).
 				Return([]byte(`{"data":[]}`), nil)
 
 			err := subject.GenerateImage(inputText, outputFile)
@@ -1078,7 +1080,7 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 		})
 		it("throws an error when base64 is invalid", func() {
 			mockCaller.EXPECT().
-				Post(subject.Config.URL+subject.Config.DrawPath, body, false).
+				Post(subject.Config.URL+subject.Config.ImageGenerationsPath, body, false).
 				Return([]byte(`{"data":[{"b64_json":"!!notbase64!!"}]}`), nil)
 
 			err := subject.GenerateImage(inputText, outputFile)
@@ -1089,7 +1091,7 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 			valid := base64.StdEncoding.EncodeToString([]byte("image-bytes"))
 
 			mockCaller.EXPECT().
-				Post(subject.Config.URL+subject.Config.DrawPath, body, false).
+				Post(subject.Config.URL+subject.Config.ImageGenerationsPath, body, false).
 				Return([]byte(fmt.Sprintf(`{"data":[{"b64_json":"%s"}]}`, valid)), nil)
 
 			mockWriter.EXPECT().Create(outputFile).Return(nil, errors.New(errorText))
@@ -1105,7 +1107,7 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 			defer file.Close()
 
 			mockCaller.EXPECT().
-				Post(subject.Config.URL+subject.Config.DrawPath, body, false).
+				Post(subject.Config.URL+subject.Config.ImageGenerationsPath, body, false).
 				Return([]byte(fmt.Sprintf(`{"data":[{"b64_json":"%s"}]}`, valid)), nil)
 
 			mockWriter.EXPECT().Create(outputFile).Return(file, nil)
@@ -1122,7 +1124,7 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 			defer file.Close()
 
 			mockCaller.EXPECT().
-				Post(subject.Config.URL+subject.Config.DrawPath, body, false).
+				Post(subject.Config.URL+subject.Config.ImageGenerationsPath, body, false).
 				Return([]byte(fmt.Sprintf(`{"data":[{"b64_json":"%s"}]}`, valid)), nil)
 
 			mockWriter.EXPECT().Create(outputFile).Return(file, nil)
@@ -1132,6 +1134,100 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 			Expect(err).NotTo(HaveOccurred())
 		})
 	})
+	when("EditImage()", func() {
+		const (
+			inputText  = "give the dog sunglasses"
+			inputFile  = "dog.png"
+			outputFile = "dog_cool.png"
+			errorText  = "mock error occurred"
+		)
+
+		var (
+			subject    *client.Client
+			validB64   string
+			imageBytes = []byte("image-bytes")
+			respBytes  []byte
+		)
+
+		it.Before(func() {
+			subject = factory.buildClientWithoutConfig()
+			validB64 = base64.StdEncoding.EncodeToString(imageBytes)
+			respBytes = []byte(fmt.Sprintf(`{"data":[{"b64_json":"%s"}]}`, validB64))
+		})
+
+		it("returns error when input file can't be opened", func() {
+			mockReader.EXPECT().Open(inputFile).Return(nil, errors.New(errorText))
+
+			err := subject.EditImage(inputText, inputFile, outputFile)
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("failed to open input image"))
+		})
+		it("returns error on invalid mime type", func() {
+			file := openDummy()
+			mockReader.EXPECT().Open(inputFile).Return(file, nil).Times(2)
+			mockReader.EXPECT().ReadBufferFromFile(file).Return([]byte("not an image"), nil)
+
+			err := subject.EditImage(inputText, inputFile, outputFile)
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("unsupported MIME type"))
+		})
+		it("returns error when HTTP call fails", func() {
+			mockReader.EXPECT().Open(inputFile).DoAndReturn(func(string) (*os.File, error) {
+				return openDummy(), nil
+			}).Times(2)
+
+			mockReader.EXPECT().
+				ReadBufferFromFile(gomock.AssignableToTypeOf(&os.File{})).
+				Return([]byte("\x89PNG\r\n\x1a\n"), nil)
+
+			mockCaller.EXPECT().
+				PostWithHeaders(gomock.Any(), gomock.Any(), gomock.Any()).
+				Return(nil, errors.New(errorText))
+
+			err := subject.EditImage(inputText, inputFile, outputFile)
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("failed to edit image"))
+		})
+		it("returns error when base64 is invalid", func() {
+			invalidResp := []byte(`{"data":[{"b64_json":"!notbase64"}]}`)
+
+			mockReader.EXPECT().Open(inputFile).DoAndReturn(func(string) (*os.File, error) {
+				return openDummy(), nil
+			}).Times(2)
+
+			mockReader.EXPECT().
+				ReadBufferFromFile(gomock.AssignableToTypeOf(&os.File{})).
+				Return([]byte("\x89PNG\r\n\x1a\n"), nil)
+
+			mockCaller.EXPECT().
+				PostWithHeaders(gomock.Any(), gomock.Any(), gomock.Any()).
+				Return(invalidResp, nil)
+
+			err := subject.EditImage(inputText, inputFile, outputFile)
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("failed to decode base64 image"))
+		})
+		it("writes image when all steps succeed", func() {
+			file := openDummy()
+			mockReader.EXPECT().Open(inputFile).DoAndReturn(func(string) (*os.File, error) {
+				return openDummy(), nil
+			}).Times(2)
+
+			mockReader.EXPECT().
+				ReadBufferFromFile(gomock.AssignableToTypeOf(&os.File{})).
+				Return([]byte("\x89PNG\r\n\x1a\n"), nil)
+
+			mockCaller.EXPECT().
+				PostWithHeaders(gomock.Any(), gomock.Any(), gomock.Any()).
+				Return(respBytes, nil)
+
+			mockWriter.EXPECT().Create(outputFile).Return(file, nil)
+			mockWriter.EXPECT().Write(file, imageBytes).Return(nil)
+
+			err := subject.EditImage(inputText, inputFile, outputFile)
+			Expect(err).NotTo(HaveOccurred())
+		})
+	})
 	when("Transcribe()", func() {
 		const audioPath = "path/to/audio.wav"
 		const transcribedText = "Hello, this is a test."
@@ -1535,6 +1631,16 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 	})
 }
 
+func openDummy() *os.File {
+	// Use os.Pipe to get an *os.File without needing a real disk file.
+	r, w, _ := os.Pipe()
+	go func() {
+		_, _ = io.Copy(w, bytes.NewBuffer([]byte("\x89PNG\r\n\x1a\n")))
+		_ = w.Close()
+	}()
+	return r
+}
+
 func createBody(messages []api.Message, stream bool) ([]byte, error) {
 	req := api.CompletionsRequest{
 		Model:            config.Model,