diff --git a/server/internal/infrastructure/piagent/image_generation.go b/server/internal/infrastructure/piagent/image_generation.go index 2996274..c826fd1 100644 --- a/server/internal/infrastructure/piagent/image_generation.go +++ b/server/internal/infrastructure/piagent/image_generation.go @@ -21,6 +21,7 @@ import ( "unicode/utf8" "img_infinite_canvas/internal/domain/design" + "img_infinite_canvas/internal/infrastructure/netutil" "github.com/zeromicro/go-zero/core/logx" ) @@ -219,7 +220,7 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) { if len(images) > 0 { - if c.inputImageTransport == "url" { + if c.inputImageTransport == "url" && imageInputsCanUseURLTransport(images) { return jsonImageEditBody(model, prompt, size, count, images, stream) } return c.multipartImageEditBody(ctx, model, prompt, size, count, images, stream) @@ -242,6 +243,38 @@ func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt return body, "application/json", err } +func imageInputsCanUseURLTransport(images []string) bool { + for _, image := range images { + if imageInputRequiresFileUpload(image) { + return false + } + } + return true +} + +func imageInputRequiresFileUpload(value string) bool { + value = strings.TrimSpace(value) + if value == "" { + return false + } + if strings.HasPrefix(value, "data:image/") || strings.HasPrefix(value, "/") { + return true + } + parsed, err := url.Parse(value) + if err != nil { + return true + } + scheme := strings.ToLower(parsed.Scheme) + if scheme != "http" && scheme != "https" { + return true + } + host := strings.ToLower(parsed.Hostname()) + if host == "" { + return true + } + return netutil.IsLocalOrPrivateHost(host) +} + func jsonImageEditBody(model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) { requestPayload := map[string]any{ "model": model, diff --git a/server/internal/infrastructure/piagent/image_generation_test.go b/server/internal/infrastructure/piagent/image_generation_test.go index a9b392f..8c38979 100644 --- a/server/internal/infrastructure/piagent/image_generation_test.go +++ b/server/internal/infrastructure/piagent/image_generation_test.go @@ -234,6 +234,60 @@ func TestImageClientSendsReferenceImagesByURLWhenConfigured(t *testing.T) { } } +func TestImageClientUploadsLocalhostReferenceImagesEvenWhenURLTransportConfigured(t *testing.T) { + var sawMultipartEdit bool + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/reference.png": + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write([]byte("local image bytes")) + case "/v1/images/edits": + sawMultipartEdit = true + if !strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data") { + t.Errorf("expected multipart edit request for local image, got %s", r.Header.Get("Content-Type")) + } + if err := r.ParseMultipartForm(32 << 20); err != nil { + t.Errorf("parse multipart form: %v", err) + return + } + file, header, err := r.FormFile("image") + if err != nil { + t.Errorf("expected image file field: %v", err) + return + } + defer file.Close() + data, err := io.ReadAll(file) + if err != nil { + t.Errorf("read image file: %v", err) + return + } + if header.Filename != "reference.png" || header.Header.Get("Content-Type") != "image/png" || string(data) != "local image bytes" { + t.Errorf("unexpected uploaded image: filename=%q content_type=%q data=%q", header.Filename, header.Header.Get("Content-Type"), string(data)) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/edited.png"}]}`)) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key", InputImageTransport: "url"}) + image, err := client.Generate(context.Background(), "edit the local image", ImageRequestOptions{ + Model: "gpt-image-2", + Images: []string{server.URL + "/reference.png"}, + }) + if err != nil { + t.Fatal(err) + } + if image.URL != "https://example.com/edited.png" { + t.Fatalf("expected edited image, got %#v", image) + } + if !sawMultipartEdit { + t.Fatal("expected multipart image edit request") + } +} + func TestImageClientUploadsInternalReferenceImagesAsMultipart(t *testing.T) { var sawEdit bool server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {