feat(image): pass reference images to GPT image model for edits

Server: forward canvas/reference images to the image model so extract, continue, and edit requests stay faithful to the source product. - Add ImageRequestOptions.Images and route requests with inputs to the /v1/images/edits endpoint, via multipart upload (default) or JSON image_url payload, selectable through Agent.Image.InputImageTransport. - Extract reference URLs from prompt directives, inline @image tokens, and chat message image/mask contents; fold them into the idempotency key and prepend a reference-fidelity instruction to the prompt. Frontend: - Shrink CanvasToast to a compact top pill. - Paste an internal transparent clipboard node instead of re-uploading when it matches the clipboard image. - Keep agent-project polling alive on stream error while streaming.
2026-07-08 01:19:20 +08:00
parent bb7fce9171
commit 2f5291a0f9
11 changed files with 614 additions and 63 deletions
@@ -3,11 +3,18 @@ package piagent
 import (
 	"bytes"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
+	"mime/multipart"
 	"net"
 	"net/http"
+	"net/textproto"
+	"net/url"
+	"path/filepath"
+	"strconv"
 	"strings"
 	"time"
 	"unicode/utf8"
@@ -35,19 +42,22 @@ type ImageRequestOptions struct {
 	Model          string
 	Size           string
 	Count          int
+	Images         []string
 	IdempotencyKey string
 }

 type ImageClientOptions struct {
-	BaseURL        string
-	APIKey         string
-	TimeoutSeconds int
+	BaseURL             string
+	APIKey              string
+	InputImageTransport string
+	TimeoutSeconds      int
 }

 type ImageClient struct {
-	baseURL string
-	apiKey  string
-	client  *http.Client
+	baseURL             string
+	apiKey              string
+	inputImageTransport string
+	client              *http.Client
 }

 func NewImageClient(opts ImageClientOptions) *ImageClient {
@@ -63,8 +73,9 @@ func NewImageClient(opts ImageClientOptions) *ImageClient {
 	}

 	return &ImageClient{
-		baseURL: strings.TrimRight(opts.BaseURL, "/"),
-		apiKey:  opts.APIKey,
+		baseURL:             strings.TrimRight(opts.BaseURL, "/"),
+		apiKey:              opts.APIKey,
+		inputImageTransport: normalizeInputImageTransport(opts.InputImageTransport),
 		client: &http.Client{
 			Timeout: timeout,
 			Transport: &http.Transport{
@@ -79,6 +90,15 @@ func NewImageClient(opts ImageClientOptions) *ImageClient {
 	}
 }

+func normalizeInputImageTransport(value string) string {
+	switch strings.ToLower(strings.TrimSpace(value)) {
+	case "url", "image_url", "image-url":
+		return "url"
+	default:
+		return "file"
+	}
+}
+
 func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error) {
 	prompt = strings.TrimSpace(prompt)
 	if prompt == "" {
@@ -99,20 +119,10 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq
 	if count > 10 {
 		count = 10
 	}
+	images := normalizeImageInputs(opts.Images)

-	requestPayload := map[string]any{
-		"model":  model,
-		"prompt": prompt,
-	}
-	if count > 1 {
-		requestPayload["n"] = count
-	}
-	if size != "" {
-		requestPayload["size"] = size
-	}
 	idempotencyKey := strings.TrimSpace(opts.IdempotencyKey)
-
-	body, err := json.Marshal(requestPayload)
+	body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, count, images)
 	if err != nil {
 		return GeneratedImage{}, err
 	}
@@ -121,7 +131,7 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq
 	startedAt := time.Now()
 	promptRunes := utf8.RuneCountInString(prompt)
 	for attempt := 0; attempt < 3; attempt++ {
-		image, err := c.generateOnce(ctx, body, idempotencyKey)
+		image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0)
 		if err == nil {
 			logx.Infof(
 				"gpt image generation succeeded model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s timeout=%s images=%d idempotency_key=%s",
@@ -166,13 +176,240 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq
 	return GeneratedImage{}, lastErr
 }

-func (c *ImageClient) generateOnce(ctx context.Context, body []byte, idempotencyKey string) (GeneratedImage, error) {
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/images/generations", bytes.NewReader(body))
+func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, count int, images []string) ([]byte, string, error) {
+	if len(images) > 0 {
+		if c.inputImageTransport == "url" {
+			return jsonImageEditBody(model, prompt, size, count, images)
+		}
+		return c.multipartImageEditBody(ctx, model, prompt, size, count, images)
+	}
+
+	requestPayload := map[string]any{
+		"model":  model,
+		"prompt": prompt,
+	}
+	if count > 1 {
+		requestPayload["n"] = count
+	}
+	if size != "" {
+		requestPayload["size"] = size
+	}
+	body, err := json.Marshal(requestPayload)
+	return body, "application/json", err
+}
+
+func jsonImageEditBody(model string, prompt string, size string, count int, images []string) ([]byte, string, error) {
+	requestPayload := map[string]any{
+		"model":  model,
+		"prompt": prompt,
+		"images": imageReferencePayload(images),
+	}
+	if count > 1 {
+		requestPayload["n"] = count
+	}
+	if size != "" {
+		requestPayload["size"] = size
+	}
+	body, err := json.Marshal(requestPayload)
+	return body, "application/json", err
+}
+
+func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, prompt string, size string, count int, images []string) ([]byte, string, error) {
+	var body bytes.Buffer
+	writer := multipart.NewWriter(&body)
+	fields := map[string]string{
+		"model":  model,
+		"prompt": prompt,
+	}
+	if size != "" {
+		fields["size"] = size
+	}
+	if count > 1 {
+		fields["n"] = strconv.Itoa(count)
+	}
+	for key, value := range fields {
+		if err := writer.WriteField(key, value); err != nil {
+			return nil, "", err
+		}
+	}
+
+	fieldName := "image"
+	if len(images) > 1 {
+		fieldName = "image[]"
+	}
+	for index, image := range images {
+		input, err := c.imageInputFile(ctx, image, index)
+		if err != nil {
+			_ = writer.Close()
+			return nil, "", err
+		}
+		part, err := createImageFormFile(writer, fieldName, input)
+		if err != nil {
+			_ = writer.Close()
+			return nil, "", err
+		}
+		if _, err := part.Write(input.Data); err != nil {
+			_ = writer.Close()
+			return nil, "", err
+		}
+	}
+	if err := writer.Close(); err != nil {
+		return nil, "", err
+	}
+	return body.Bytes(), writer.FormDataContentType(), nil
+}
+
+func createImageFormFile(writer *multipart.Writer, fieldName string, input imageInputFile) (io.Writer, error) {
+	header := make(textproto.MIMEHeader)
+	header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, escapeMultipartValue(fieldName), escapeMultipartValue(input.FileName)))
+	contentType := strings.TrimSpace(input.ContentType)
+	if contentType == "" {
+		contentType = "application/octet-stream"
+	}
+	header.Set("Content-Type", contentType)
+	return writer.CreatePart(header)
+}
+
+func escapeMultipartValue(value string) string {
+	return strings.NewReplacer("\\", "\\\\", `"`, "\\\"").Replace(value)
+}
+
+type imageInputFile struct {
+	FileName    string
+	ContentType string
+	Data        []byte
+}
+
+func (c *ImageClient) imageInputFile(ctx context.Context, value string, index int) (imageInputFile, error) {
+	value = strings.TrimSpace(value)
+	if strings.HasPrefix(value, "data:image/") {
+		header, payload, ok := strings.Cut(value, ",")
+		if !ok {
+			return imageInputFile{}, fmt.Errorf("invalid image data url")
+		}
+		contentType := strings.TrimPrefix(strings.Split(header, ";")[0], "data:")
+		data, err := base64.StdEncoding.DecodeString(payload)
+		if err != nil {
+			return imageInputFile{}, err
+		}
+		return imageInputFile{
+			FileName:    fmt.Sprintf("reference-%d%s", index+1, imageExtension(contentType, "")),
+			ContentType: contentType,
+			Data:        data,
+		}, nil
+	}
+	if !strings.HasPrefix(value, "http://") && !strings.HasPrefix(value, "https://") {
+		return imageInputFile{}, fmt.Errorf("unsupported image input %q", value)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, value, nil)
+	if err != nil {
+		return imageInputFile{}, err
+	}
+	req.Header.Set("User-Agent", "img-infinite-canvas/1.0")
+	resp, err := c.client.Do(req)
+	if err != nil {
+		return imageInputFile{}, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return imageInputFile{}, fmt.Errorf("download reference image failed: %s", resp.Status)
+	}
+	contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
+	data, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20))
+	if err != nil {
+		return imageInputFile{}, err
+	}
+	if len(data) == 0 {
+		return imageInputFile{}, fmt.Errorf("download reference image returned empty body")
+	}
+	if contentType == "" || !strings.HasPrefix(contentType, "image/") {
+		contentType = http.DetectContentType(data)
+	}
+	if !strings.HasPrefix(contentType, "image/") {
+		return imageInputFile{}, fmt.Errorf("reference input is not an image: %s", contentType)
+	}
+	return imageInputFile{
+		FileName:    imageFileName(value, contentType, index),
+		ContentType: contentType,
+		Data:        data,
+	}, nil
+}
+
+func normalizeImageInputs(images []string) []string {
+	if len(images) == 0 {
+		return nil
+	}
+	normalized := make([]string, 0, len(images))
+	seen := make(map[string]struct{})
+	for _, image := range images {
+		image = strings.TrimSpace(image)
+		if image == "" {
+			continue
+		}
+		if _, ok := seen[image]; ok {
+			continue
+		}
+		seen[image] = struct{}{}
+		normalized = append(normalized, image)
+		if len(normalized) >= 16 {
+			break
+		}
+	}
+	return normalized
+}
+
+func imageReferencePayload(images []string) []map[string]string {
+	items := make([]map[string]string, 0, len(images))
+	for _, image := range images {
+		items = append(items, map[string]string{"image_url": image})
+	}
+	return items
+}
+
+func imageFileName(value string, contentType string, index int) string {
+	name := ""
+	if parsed, err := url.Parse(value); err == nil {
+		name = filepath.Base(parsed.Path)
+	}
+	if name == "" || name == "." || name == "/" {
+		name = fmt.Sprintf("reference-%d%s", index+1, imageExtension(contentType, ""))
+	}
+	if filepath.Ext(name) == "" {
+		name += imageExtension(contentType, ".png")
+	}
+	return name
+}
+
+func imageExtension(contentType string, fallback string) string {
+	switch strings.ToLower(strings.TrimSpace(contentType)) {
+	case "image/png":
+		return ".png"
+	case "image/jpeg", "image/jpg":
+		return ".jpg"
+	case "image/webp":
+		return ".webp"
+	case "image/gif":
+		return ".gif"
+	default:
+		if fallback != "" {
+			return fallback
+		}
+		return ".png"
+	}
+}
+
+func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType string, idempotencyKey string, edit bool) (GeneratedImage, error) {
+	endpoint := "/v1/images/generations"
+	if edit {
+		endpoint = "/v1/images/edits"
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+endpoint, bytes.NewReader(body))
 	if err != nil {
 		return GeneratedImage{}, err
 	}
 	req.Header.Set("Authorization", "Bearer "+c.apiKey)
-	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Content-Type", contentType)
 	req.Header.Set("Connection", "close")
 	if idempotencyKey != "" {
 		req.Header.Set("Idempotency-Key", idempotencyKey)