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.
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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{})
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user