From 2f5291a0f937ac9bd487f840279f884e163bae0b Mon Sep 17 00:00:00 2001 From: liangxu Date: Wed, 8 Jul 2026 01:19:20 +0800 Subject: [PATCH] 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. --- .../src/ui/components/CanvasToast/index.css | 56 ++-- .../canvas/hooks/useGenerationStream.ts | 8 +- .../src/ui/pages/CanvasWorkspace/index.tsx | 24 +- server/README.md | 3 + server/etc/config.yaml | 1 + server/internal/config/config.go | 9 +- .../infrastructure/piagent/design_agent.go | 74 +++++ .../piagent/design_agent_test.go | 46 +++ .../piagent/image_generation.go | 285 ++++++++++++++++-- .../piagent/image_generation_test.go | 164 ++++++++++ server/internal/svc/service_context.go | 7 +- 11 files changed, 614 insertions(+), 63 deletions(-) diff --git a/frontend/src/ui/components/CanvasToast/index.css b/frontend/src/ui/components/CanvasToast/index.css index b563772..d51cfa8 100644 --- a/frontend/src/ui/components/CanvasToast/index.css +++ b/frontend/src/ui/components/CanvasToast/index.css @@ -1,27 +1,30 @@ .canvas-toast { position: fixed; left: 50%; - top: 48px; + top: 14px; z-index: 120; - width: min(1040px, calc(100vw - 140px)); - min-height: 126px; + box-sizing: border-box; + width: fit-content; + min-width: 220px; + max-width: min(480px, calc(100vw - 32px)); + min-height: 54px; display: flex; align-items: center; - gap: 28px; - padding: 0 44px; + gap: 12px; + padding: 12px 18px; border: 1px solid rgba(224, 229, 238, 0.96); - border-radius: 24px; + border-radius: 14px; color: #1f2329; background: rgba(255, 255, 255, 0.98); - box-shadow: 0 38px 90px rgba(17, 24, 39, 0.16); + box-shadow: 0 14px 36px rgba(17, 24, 39, 0.12); pointer-events: none; transform: translateX(-50%); animation: canvas-toast-enter 180ms cubic-bezier(0.2, 0, 0.2, 1); } .canvas-toast-icon { - width: 44px; - height: 44px; + width: 30px; + height: 30px; display: grid; place-items: center; flex: 0 0 auto; @@ -29,6 +32,11 @@ color: #ffffff; } +.canvas-toast-icon svg { + width: 18px; + height: 18px; +} + .canvas-toast.success .canvas-toast-icon { background: #45d27c; } @@ -40,9 +48,9 @@ .canvas-toast strong { min-width: 0; overflow: hidden; - font-size: 32px; - font-weight: 850; - line-height: 1.18; + font-size: 15px; + font-weight: 750; + line-height: 1.32; text-overflow: ellipsis; white-space: normal; word-break: break-word; @@ -62,26 +70,28 @@ @media (max-width: 760px) { .canvas-toast { - top: 18px; - width: calc(100vw - 32px); - min-height: 76px; - gap: 14px; - padding: 16px 18px; - border-radius: 18px; + top: 0; + width: fit-content; + min-width: 0; + max-width: calc(100vw - 28px); + min-height: 50px; + gap: 10px; + padding: 11px 14px; + border-radius: 13px; } .canvas-toast-icon { - width: 32px; - height: 32px; + width: 28px; + height: 28px; } .canvas-toast-icon svg { - width: 20px; - height: 20px; + width: 17px; + height: 17px; } .canvas-toast strong { - font-size: 18px; + font-size: 14px; line-height: 1.28; } } diff --git a/frontend/src/ui/pages/CanvasWorkspace/canvas/hooks/useGenerationStream.ts b/frontend/src/ui/pages/CanvasWorkspace/canvas/hooks/useGenerationStream.ts index c3af9e8..49fc6ec 100644 --- a/frontend/src/ui/pages/CanvasWorkspace/canvas/hooks/useGenerationStream.ts +++ b/frontend/src/ui/pages/CanvasWorkspace/canvas/hooks/useGenerationStream.ts @@ -233,8 +233,9 @@ export function useGenerationStream({ closeRef.current = null; }, onError: (nextError) => { - if (closeAgentProjectPollRef.current) { + if (streamingRef.current === streamKey) { closeRef.current = null; + startAgentProjectPolling(nextProjectId, threadId); return; } completeChannel(channel); @@ -243,7 +244,6 @@ export function useGenerationStream({ closeRef.current = null; } }); - startAgentProjectPolling(nextProjectId, threadId); }, [appendThinkingMessage, applyProject, completeChannel, toUserMessageText, messages.generateError, messages.generateTimeout, setError, startAgentProjectPolling, startGeneratorTaskPolling] ); @@ -306,8 +306,9 @@ export function useGenerationStream({ closeRef.current = null; }, onError: (nextError) => { - if (closeAgentProjectPollRef.current) { + if (streamingRef.current === streamKey) { closeRef.current = null; + startAgentProjectPolling(nextProjectId, payload.threadId); return; } completeChannel(channel); @@ -316,7 +317,6 @@ export function useGenerationStream({ closeRef.current = null; } }); - startAgentProjectPolling(nextProjectId, payload.threadId); }, [appendThinkingMessage, applyProject, completeChannel, toUserMessageText, messages.generateError, messages.generateTimeout, setError, startAgentProjectPolling, startGeneratorTaskPolling] ); diff --git a/frontend/src/ui/pages/CanvasWorkspace/index.tsx b/frontend/src/ui/pages/CanvasWorkspace/index.tsx index 8d73346..c6f8bc8 100644 --- a/frontend/src/ui/pages/CanvasWorkspace/index.tsx +++ b/frontend/src/ui/pages/CanvasWorkspace/index.tsx @@ -1843,9 +1843,9 @@ export function CanvasWorkspace({ projectId }: { projectId: string }) { .catch(() => showCanvasToast(t("clipboardCopyFailed"), "error")); }; - const pasteNode = () => { - if (!clipboardNode) return; - const copy = cloneCanvasNode(clipboardNode, cryptoFallbackId()); + const pasteNode = (sourceNode = clipboardNode) => { + if (!sourceNode) return; + const copy = cloneCanvasNode(sourceNode, cryptoFallbackId()); setNodes((current) => [...current, copy]); selectOnlyNode(copy.id); setDirty(true); @@ -2145,11 +2145,25 @@ export function CanvasWorkspace({ projectId }: { projectId: string }) { if (isEditableShortcutTarget(event.target)) return; const imageFiles = event.clipboardData ? dataTransferImageFiles(event.clipboardData, "clipboard-image") : []; - if (imageFiles.length > 0) { - event.preventDefault(); + const pasteImageFiles = () => { const pointer = lastCanvasPointerRef.current; const anchor = pointer ? canvasPointFromClient(pointer.clientX, pointer.clientY) : canvasCenterPoint(); void uploadImagesToCanvas(imageFiles, anchor); + }; + if (clipboardNode && imageFiles.length > 0 && isTransparentClipboardNode(clipboardNode)) { + event.preventDefault(); + void shouldUseInternalClipboardNode(clipboardNode, imageFiles).then((useInternalNode) => { + if (useInternalNode) { + pasteNode(clipboardNode); + return; + } + pasteImageFiles(); + }); + return; + } + if (imageFiles.length > 0) { + event.preventDefault(); + pasteImageFiles(); return; } diff --git a/server/README.md b/server/README.md index 706d417..5c6fc56 100644 --- a/server/README.md +++ b/server/README.md @@ -58,6 +58,7 @@ Agent: Image: BaseURL: http://154.21.87.3:8080 APIKeyEnv: GPT_IMAGE_API_KEY + InputImageTransport: file TimeoutSeconds: 600 ObjectStorage: Driver: minio @@ -71,6 +72,8 @@ ObjectStorage: ExpiresIn: 900 ``` +`Agent.Image.InputImageTransport` controls how input/reference images are sent to the image model for edit/image-to-image calls: `file` uploads image bytes with multipart form data, while `url` sends `images: [{"image_url": "..."}]` for gateways that can fetch public URLs directly. + Start local Docker dependencies for debugging: ```bash diff --git a/server/etc/config.yaml b/server/etc/config.yaml index 1c8cf93..819c5f0 100644 --- a/server/etc/config.yaml +++ b/server/etc/config.yaml @@ -104,6 +104,7 @@ Agent: BaseURL: http://154.21.87.3:8080 APIKeyEnv: GPT_IMAGE_API_KEY APIKey: "sk-0af1835a4e855198ec6e1ff8fb120844c146dd09b0b25b05b30ed8fa5c7be568" + InputImageTransport: file TimeoutSeconds: 600 ObjectStorage: Driver: minio diff --git a/server/internal/config/config.go b/server/internal/config/config.go index 0c906b0..af4735b 100644 --- a/server/internal/config/config.go +++ b/server/internal/config/config.go @@ -123,10 +123,11 @@ type AgentConfig struct { } type ImageGenerationConfig struct { - BaseURL string `json:",optional"` - APIKey string `json:",optional"` - APIKeyEnv string `json:",default=GPT_IMAGE_API_KEY"` - TimeoutSeconds int `json:",default=600"` + BaseURL string `json:",optional"` + APIKey string `json:",optional"` + APIKeyEnv string `json:",default=GPT_IMAGE_API_KEY"` + InputImageTransport string `json:",default=file"` + TimeoutSeconds int `json:",default=600"` } type TextExtractionConfig struct { diff --git a/server/internal/infrastructure/piagent/design_agent.go b/server/internal/infrastructure/piagent/design_agent.go index e46f46c..740f25d 100644 --- a/server/internal/infrastructure/piagent/design_agent.go +++ b/server/internal/infrastructure/piagent/design_agent.go @@ -283,6 +283,9 @@ func (a *DesignAgent) generatePlannedImagesIncremental(ctx context.Context, prom } func (a *DesignAgent) generateImageWithFallback(ctx context.Context, req design.AgentRequest, prompt string, fallback string, index int) (GeneratedImage, error) { + referenceImages := referenceImageURLs(req) + prompt = withReferenceImageInstruction(prompt, len(referenceImages)) + fallback = withReferenceImageInstruction(fallback, len(referenceImages)) prompt = prepareGPTImagePrompt(prompt, imageGenerationPromptMaxRunes) fallback = prepareGPTImagePrompt(fallback, imageGenerationFallbackMaxRunes) if imagePromptTooLarge(prompt) && strings.TrimSpace(fallback) != "" { @@ -293,6 +296,7 @@ func (a *DesignAgent) generateImageWithFallback(ctx context.Context, req design. Model: req.ImageModel, Size: req.ImageSize, Count: 1, + Images: referenceImages, IdempotencyKey: imageGenerationIdempotencyKey(req, index, prompt, false), }) if err == nil { @@ -308,6 +312,7 @@ func (a *DesignAgent) generateImageWithFallback(ctx context.Context, req design. Model: req.ImageModel, Size: req.ImageSize, Count: 1, + Images: referenceImages, IdempotencyKey: imageGenerationIdempotencyKey(req, index, fallback, true), }) if fallbackErr == nil { @@ -329,12 +334,79 @@ func imageGenerationIdempotencyKey(req design.AgentRequest, index int, prompt st strings.TrimSpace(req.ImageSize), strconv.Itoa(index), kind, + strings.Join(referenceImageURLs(req), "\n"), strings.TrimSpace(prompt), }, "\n") sum := sha256.Sum256([]byte(seed)) return "img-gen-" + kind + "-" + hex.EncodeToString(sum[:16]) } +func referenceImageURLs(req design.AgentRequest) []string { + urls := make([]string, 0) + seen := make(map[string]struct{}) + add := func(value string) { + value = strings.TrimSpace(strings.TrimRight(value, "，。,.")) + if value == "" { + return + } + if _, ok := seen[value]; ok { + return + } + seen[value] = struct{}{} + urls = append(urls, value) + } + addFromText := func(text string) { + for _, match := range referenceImageDirectiveURLPattern.FindAllStringSubmatch(text, -1) { + if len(match) > 1 { + add(match[1]) + } + } + for _, match := range promptImageTokenURLPattern.FindAllStringSubmatch(text, -1) { + if len(match) > 1 && strings.TrimSpace(match[1]) != "" { + add(match[1]) + continue + } + if len(match) > 2 { + add(match[2]) + } + } + } + + addFromText(req.Prompt) + for _, message := range req.Messages { + for _, content := range message.Contents { + switch strings.ToLower(strings.TrimSpace(content.Type)) { + case "image", "mask": + add(content.ImageURL) + case "text": + addFromText(content.Text) + } + if len(urls) >= 16 { + return urls[:16] + } + } + } + if len(urls) > 16 { + return urls[:16] + } + return urls +} + +func withReferenceImageInstruction(prompt string, imageCount int) string { + prompt = strings.TrimSpace(prompt) + if prompt == "" || imageCount <= 0 { + return prompt + } + lower := strings.ToLower(prompt) + if strings.Contains(prompt, "附带参考图") || strings.Contains(lower, "attached reference image") || strings.Contains(lower, "input reference image") { + return prompt + } + if prefersEnglish(prompt) { + return "Input reference image(s) are attached. Use them as the visual source of truth for product shape, fabric, color, structure, and identity; when extracting, continuing, or editing a product, do not invent a different item.\n" + prompt + } + return "已随请求附带参考图；生成时必须以参考图中的商品外形、材质、颜色、结构细节和身份为准，拆出、延续或编辑商品时不要另造不同商品。\n" + prompt +} + func imagePromptTooLarge(prompt string) bool { return utf8.RuneCountInString(strings.TrimSpace(prompt)) > imageGenerationPromptTooLargeRunes } @@ -404,6 +476,8 @@ var explicitPixelSizePattern = regexp.MustCompile(`(?i)([1-9]\d{2,4})\s*(?:x|×| var explicitAspectRatioPattern = regexp.MustCompile(`(?i)(?:^|[^\d])((?:1|2|3|4|9|16)\s*:\s*(?:1|2|3|4|9|16))(?:[^\d]|$)`) var imagePromptURLPattern = regexp.MustCompile(`https?://\S+`) +var referenceImageDirectiveURLPattern = regexp.MustCompile(`(?i)(?:参考图|reference image)\s*\d*\s*(?:[（(][^）)\n]+[）)])?\s*[:：]\s*(https?://[^\s\]]+|data:image/[^\s\]]+|/[^\s\]]+)`) +var promptImageTokenURLPattern = regexp.MustCompile(`\[@image:#\d+:[^:\]]+:(?:\[[^\]]+\]\(([^)]+)\)|([^\]]+))\]`) func fullFrameImagePrompt(prompt string) string { prompt = strings.TrimSpace(prompt) diff --git a/server/internal/infrastructure/piagent/design_agent_test.go b/server/internal/infrastructure/piagent/design_agent_test.go index 0a69950..5ea3da4 100644 --- a/server/internal/infrastructure/piagent/design_agent_test.go +++ b/server/internal/infrastructure/piagent/design_agent_test.go @@ -214,6 +214,52 @@ func TestDesignAgentGeneratesImageThroughPiAgent(t *testing.T) { } } +func TestDesignAgentPassesInlineReferenceImagesToImageGenerator(t *testing.T) { + imageGenerator := &fakeImageGenerator{} + agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{}) + + _, err := agent.Plan(context.Background(), design.AgentRequest{ + Prompt: "把参考图里的灰裤子拆成单件商品图", + Mode: "image", + ImageModel: "gpt-image-2", + Messages: []design.AgentChatMessage{ + { + Role: "user", + Contents: []design.AgentContent{ + {Type: "image", ImageURL: "https://example.com/pants-rack.png", ImageWidth: 1200, ImageHeight: 900}, + {Type: "text", Text: "把参考图里的灰裤子拆成单件商品图", TextSource: "input"}, + }, + }, + }, + }) + if err != nil { + t.Fatal(err) + } + if len(imageGenerator.opts.Images) != 1 || imageGenerator.opts.Images[0] != "https://example.com/pants-rack.png" { + t.Fatalf("expected reference image to be passed to GPT Image 2, got %#v", imageGenerator.opts.Images) + } + if !strings.Contains(imageGenerator.prompt, "参考图") { + t.Fatalf("expected prompt to tell image model to use attached reference, got %s", imageGenerator.prompt) + } +} + +func TestDesignAgentPassesPromptDirectiveReferenceImagesToImageGenerator(t *testing.T) { + imageGenerator := &fakeImageGenerator{} + agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{}) + + _, err := agent.Plan(context.Background(), design.AgentRequest{ + Prompt: "参考图 1（image）：https://example.com/pants-rack.png\n灰裤子拆成单件衣服，为了后面电商配图。", + Mode: "image", + ImageModel: "gpt-image-2", + }) + if err != nil { + t.Fatal(err) + } + if len(imageGenerator.opts.Images) != 1 || imageGenerator.opts.Images[0] != "https://example.com/pants-rack.png" { + t.Fatalf("expected directive reference image to be passed to GPT Image 2, got %#v", imageGenerator.opts.Images) + } +} + func TestDesignAgentUsesDefaultImageSizeWithoutExplicitRequest(t *testing.T) { imageGenerator := &fakeImageGenerator{} agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{}) diff --git a/server/internal/infrastructure/piagent/image_generation.go b/server/internal/infrastructure/piagent/image_generation.go index ad773b7..1ee3171 100644 --- a/server/internal/infrastructure/piagent/image_generation.go +++ b/server/internal/infrastructure/piagent/image_generation.go @@ -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) diff --git a/server/internal/infrastructure/piagent/image_generation_test.go b/server/internal/infrastructure/piagent/image_generation_test.go index ee0d616..17c5cb3 100644 --- a/server/internal/infrastructure/piagent/image_generation_test.go +++ b/server/internal/infrastructure/piagent/image_generation_test.go @@ -3,8 +3,10 @@ package piagent import ( "context" "encoding/json" + "io" "net/http" "net/http/httptest" + "strings" "testing" "time" ) @@ -64,6 +66,19 @@ func TestNewImageClientDisablesKeepAliveReuse(t *testing.T) { } } +func TestNewImageClientDefaultsInputImagesToFileTransport(t *testing.T) { + client := NewImageClient(ImageClientOptions{ + BaseURL: "http://example.com", + APIKey: "test-key", + }) + if client == nil { + t.Fatal("expected image client") + } + if client.inputImageTransport != "file" { + t.Fatalf("expected file input image transport, got %q", client.inputImageTransport) + } +} + func TestImageClientSendsTransportDefaultSizeWhenUnset(t *testing.T) { var requestBody map[string]any server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -125,6 +140,155 @@ func TestImageClientSendsExplicitSize(t *testing.T) { } } +func TestImageClientUploadsDataReferenceImagesAsMultipart(t *testing.T) { + var requestPath string + var requestContentType string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestPath = r.URL.Path + requestContentType = 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-1.png" || header.Header.Get("Content-Type") != "image/png" || string(data) != "reference 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/image.png"}]}`)) + })) + defer server.Close() + + client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"}) + _, err := client.Generate(context.Background(), "extract product from reference", ImageRequestOptions{ + Model: "gpt-image-2", + Images: []string{"data:image/png;base64,cmVmZXJlbmNlIGJ5dGVz"}, + }) + if err != nil { + t.Fatal(err) + } + if requestPath != "/v1/images/edits" { + t.Fatalf("expected image edit endpoint for references, got %s", requestPath) + } + if !strings.HasPrefix(requestContentType, "multipart/form-data") { + t.Fatalf("expected multipart edit request, got %s", requestContentType) + } +} + +func TestImageClientSendsReferenceImagesByURLWhenConfigured(t *testing.T) { + var requestBody map[string]any + var requestPath string + var requestContentType string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestPath = r.URL.Path + requestContentType = r.Header.Get("Content-Type") + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Fatal(err) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`)) + })) + defer server.Close() + + client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key", InputImageTransport: "url"}) + _, err := client.Generate(context.Background(), "extract product from reference", ImageRequestOptions{ + Model: "gpt-image-2", + Images: []string{" https://example.com/reference.png ", "", "https://example.com/reference.png", "https://example.com/second.webp"}, + }) + if err != nil { + t.Fatal(err) + } + if requestPath != "/v1/images/edits" { + t.Fatalf("expected image edit endpoint for references, got %s", requestPath) + } + if requestContentType != "application/json" { + t.Fatalf("expected url transport to use JSON, got %s", requestContentType) + } + images, ok := requestBody["images"].([]any) + if !ok { + t.Fatalf("expected images array in request, got %#v", requestBody["images"]) + } + if len(images) != 2 { + t.Fatalf("expected normalized reference images, got %#v", images) + } + first, ok := images[0].(map[string]any) + if !ok || first["image_url"] != "https://example.com/reference.png" { + t.Fatalf("expected first image_url reference, got %#v", images[0]) + } + second, ok := images[1].(map[string]any) + if !ok || second["image_url"] != "https://example.com/second.webp" { + t.Fatalf("expected second image_url reference, got %#v", images[1]) + } +} + +func TestImageClientUploadsInternalReferenceImagesAsMultipart(t *testing.T) { + var sawEdit 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("internal image bytes")) + case "/v1/images/edits": + sawEdit = true + if !strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data") { + t.Errorf("expected multipart edit request, got %s", r.Header.Get("Content-Type")) + } + if err := r.ParseMultipartForm(32 << 20); err != nil { + t.Errorf("parse multipart form: %v", err) + return + } + if r.FormValue("model") != "gpt-image-2" || r.FormValue("prompt") != "edit the internal image" || r.FormValue("size") != "1024x1024" { + t.Errorf("unexpected multipart fields: model=%q prompt=%q size=%q", r.FormValue("model"), r.FormValue("prompt"), r.FormValue("size")) + } + 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) != "internal 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"}) + image, err := client.Generate(context.Background(), "edit the internal image", ImageRequestOptions{ + Model: "gpt-image-2", + Size: "1024x1024", + 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 !sawEdit { + t.Fatal("expected image edit request") + } +} + func TestImageClientParsesAlternateImageResponseShapes(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") diff --git a/server/internal/svc/service_context.go b/server/internal/svc/service_context.go index 843fac0..fdf6ff6 100644 --- a/server/internal/svc/service_context.go +++ b/server/internal/svc/service_context.go @@ -304,9 +304,10 @@ func newImageGenerator(c config.Config) piagent.ImageGenerator { opts := c.Agent.Image apiKey := resolveImageAPIKey(opts) client := piagent.NewImageClient(piagent.ImageClientOptions{ - BaseURL: opts.BaseURL, - APIKey: apiKey, - TimeoutSeconds: opts.TimeoutSeconds, + BaseURL: opts.BaseURL, + APIKey: apiKey, + InputImageTransport: opts.InputImageTransport, + TimeoutSeconds: opts.TimeoutSeconds, }) if client == nil { return nil