From cc30ae42847af0f0758fa0da3132a152488f38b9 Mon Sep 17 00:00:00 2001 From: liangxu Date: Thu, 9 Jul 2026 09:05:18 +0800 Subject: [PATCH] feat(piagent): batch-generate images in one call when generators share a prompt When multiple image-generator tasks carry the same brief, request all images from a single GPT Image call with Count=N instead of firing one call per image. Adds directImageBatchCount/Prompt detection, batch (incremental) generation paths, and folds the planned count into the idempotency key. Disables the streaming plan probe for multi-image requests. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../infrastructure/piagent/design_agent.go | 137 +++++++++++++++++- .../piagent/design_agent_test.go | 62 ++++++++ .../piagent/image_generation.go | 2 +- .../piagent/image_generation_test.go | 27 ++++ 4 files changed, 223 insertions(+), 5 deletions(-) diff --git a/server/internal/infrastructure/piagent/design_agent.go b/server/internal/infrastructure/piagent/design_agent.go index d76f9fb..8870a7d 100644 --- a/server/internal/infrastructure/piagent/design_agent.go +++ b/server/internal/infrastructure/piagent/design_agent.go @@ -129,6 +129,17 @@ func (a *DesignAgent) planDirectImageGeneration(ctx context.Context, req design. } generationReq := req generationReq.ImageSize = resolveImageSizeForGeneration(req.ImageSize, req.Prompt) + if count := directImageBatchCount(generationReq); count > 1 { + imagePrompt := directImageBatchPrompt(generationReq) + imageURLs, err := a.generateDirectImageBatch(ctx, generationReq, imagePrompt, count) + if err != nil { + return design.AgentPlan{}, fmt.Errorf("gpt image 2 generation failed: %w", err) + } + plan := directImagePlan(generationReq) + applyPlannedGeneratedImages(&plan, generationReq, imageURLs, []design.AgentImagePrompt{imagePrompt}) + plan.Message.Title = generatedMessageTitle(imagePrompt.Title, req.Prompt) + return plan, nil + } prompts := directImagePrompts(generationReq) imageURLs, err := a.generatePlannedImages(ctx, prompts, generationReq) if err != nil { @@ -150,8 +161,23 @@ func (a *DesignAgent) planDirectImageGenerationIncremental(ctx context.Context, } generationReq := req generationReq.ImageSize = resolveImageSizeForGeneration(req.ImageSize, req.Prompt) - prompts := directImagePrompts(generationReq) plan := directImagePlan(generationReq) + if count := directImageBatchCount(generationReq); count > 1 { + imagePrompt := directImageBatchPrompt(generationReq) + nodes, err := a.generateDirectImageBatchIncremental(ctx, generationReq, imagePrompt, count, plan.GeneratedNodes, onResult) + if err != nil { + if len(nodes) > 0 { + plan.GeneratedNodes = nodes + plan.Connections = nil + } + return plan, fmt.Errorf("gpt image 2 generation failed: %w", err) + } + plan.GeneratedNodes = nodes + plan.Connections = nil + plan.Message.Title = generatedMessageTitle(imagePrompt.Title, req.Prompt) + return plan, nil + } + prompts := directImagePrompts(generationReq) nodes, err := a.generatePlannedImagesIncremental(ctx, prompts, generationReq, plan.GeneratedNodes, onResult) if err != nil { if len(nodes) > 0 { @@ -204,6 +230,49 @@ func directImagePrompts(req design.AgentRequest) []design.AgentImagePrompt { return prompts } +func directImageBatchCount(req design.AgentRequest) int { + count := plannedGenerationCount(req) + if count <= 1 || !directImageTasksSharePrompt(req) { + return 0 + } + return count +} + +func directImageTasksSharePrompt(req design.AgentRequest) bool { + common := "" + for _, task := range req.TaskPlan.ImageTasks { + brief := strings.TrimSpace(task.Brief) + if brief == "" { + continue + } + if common == "" { + common = brief + continue + } + if brief != common { + return false + } + } + return true +} + +func directImageBatchPrompt(req design.AgentRequest) design.AgentImagePrompt { + prompt := strings.TrimSpace(req.Prompt) + if prompt == "" { + prompt = fallbackDirectImagePrompt(req) + } + title := normalizeGeneratedTitle("", req.Prompt) + if len(req.TaskPlan.ImageTasks) > 0 { + if taskTitle := strings.TrimSpace(req.TaskPlan.ImageTasks[0].Title); taskTitle != "" { + title = normalizeGeneratedTitle(taskTitle, req.Prompt) + } + } + return design.AgentImagePrompt{ + Title: title, + Prompt: prepareGPTImagePrompt(prompt, imageGenerationPromptMaxRunes), + } +} + func shouldUseDirectImageGeneration(req design.AgentRequest) bool { mode := strings.ToLower(strings.TrimSpace(req.Mode)) return mode == "image" || mode == "pure-image" || mode == "image-generator" || strings.Contains(mode, "image-action") @@ -295,12 +364,58 @@ func (a *DesignAgent) generatePlannedImages(ctx context.Context, prompts []desig return urls, nil } +func (a *DesignAgent) generateDirectImageBatch(ctx context.Context, req design.AgentRequest, prompt design.AgentImagePrompt, count int) ([]string, error) { + image, err := a.generateImageWithFallbackCount(ctx, req, prompt.Prompt, fallbackImagePrompt(req, 0), 0, count) + if err != nil { + return nil, err + } + urls := generatedImageURLs(image) + if len(urls) == 0 { + return nil, fmt.Errorf("image generation returned empty image") + } + if len(urls) < count { + return urls, fmt.Errorf("only %d of %d planned images returned", len(urls), count) + } + return urls[:count], nil +} + type plannedImageResult struct { index int node design.Node err error } +func (a *DesignAgent) generateDirectImageBatchIncremental( + ctx context.Context, + req design.AgentRequest, + prompt design.AgentImagePrompt, + count int, + fallback []design.Node, + onResult func(design.AgentImageGenerationResult) error, +) ([]design.Node, error) { + imageURLs, err := a.generateDirectImageBatch(ctx, req, prompt, count) + if len(imageURLs) == 0 { + return nil, err + } + slots := generatedImageSlots(req, fallback, len(imageURLs)) + nodes := make([]design.Node, 0, len(imageURLs)) + var firstErr error + for index, imageURL := range imageURLs { + node := generatedImageNodeFromSlot(slots[index], req, imageURL, prompt.Title, index, count) + if onResult != nil { + if emitErr := onResult(design.AgentImageGenerationResult{Index: index, Node: node}); emitErr != nil && firstErr == nil { + firstErr = emitErr + continue + } + } + nodes = append(nodes, node) + } + if err != nil && firstErr == nil { + firstErr = err + } + return nodes, firstErr +} + func (a *DesignAgent) generatePlannedImagesIncremental(ctx context.Context, prompts []design.AgentImagePrompt, req design.AgentRequest, fallback []design.Node, onResult func(design.AgentImageGenerationResult) error) ([]design.Node, error) { if len(prompts) == 0 { prompts = []design.AgentImagePrompt{{Title: normalizeGeneratedTitle("", req.Prompt), Prompt: prepareGPTImagePrompt(fallbackDirectImagePrompt(req), imageGenerationFallbackMaxRunes)}} @@ -377,6 +492,16 @@ 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) { + return a.generateImageWithFallbackCount(ctx, req, prompt, fallback, index, 1) +} + +func (a *DesignAgent) generateImageWithFallbackCount(ctx context.Context, req design.AgentRequest, prompt string, fallback string, index int, count int) (GeneratedImage, error) { + if count < 1 { + count = 1 + } + if count > 10 { + count = 10 + } referenceImages := referenceImageURLs(req) prompt = withReferenceImageContext(prompt, req, len(referenceImages)) fallback = withReferenceImageContext(fallback, req, len(referenceImages)) @@ -391,7 +516,7 @@ func (a *DesignAgent) generateImageWithFallback(ctx context.Context, req design. image, err := a.imageAgent.Generate(ctx, prompt, ImageRequestOptions{ Model: req.ImageModel, Size: req.ImageSize, - Count: 1, + Count: count, Images: referenceImages, IdempotencyKey: imageGenerationIdempotencyKey(req, index, prompt, false), }) @@ -407,7 +532,7 @@ func (a *DesignAgent) generateImageWithFallback(ctx context.Context, req design. fallbackImage, fallbackErr := a.imageAgent.Generate(ctx, fallback, ImageRequestOptions{ Model: req.ImageModel, Size: req.ImageSize, - Count: 1, + Count: count, Images: referenceImages, IdempotencyKey: imageGenerationIdempotencyKey(req, index, fallback, true), }) @@ -428,6 +553,7 @@ func imageGenerationIdempotencyKey(req design.AgentRequest, index int, prompt st strings.TrimSpace(req.Prompt), strings.TrimSpace(req.ImageModel), strings.TrimSpace(req.ImageSize), + strconv.Itoa(plannedGenerationCount(req)), strconv.Itoa(index), kind, strings.Join(referenceImageURLs(req), "\n"), @@ -776,7 +902,9 @@ func applyPlannedGeneratedImages(plan *design.AgentPlan, req design.AgentRequest nodes := make([]design.Node, 0, requestedCount) for index := 0; index < requestedCount; index++ { title := "" - if index < len(prompts) { + if len(prompts) == 1 { + title = prompts[0].Title + } else if index < len(prompts) { title = prompts[index].Title } nodes = append(nodes, generatedImageNodeFromSlot(slots[index], req, imageURLs[index], title, index, requestedCount)) @@ -806,6 +934,7 @@ func generatedImageNodeFromSlot(slot design.Node, req design.AgentRequest, image } if node.LayerRole == "image-generator" { node.LayerRole = "" + node.ParentID = "" } if node.Tone == "" || node.Tone == "image-generator" { node.Tone = "visual" diff --git a/server/internal/infrastructure/piagent/design_agent_test.go b/server/internal/infrastructure/piagent/design_agent_test.go index f187e00..43b9372 100644 --- a/server/internal/infrastructure/piagent/design_agent_test.go +++ b/server/internal/infrastructure/piagent/design_agent_test.go @@ -270,6 +270,64 @@ func TestDesignAgentImageGeneratorModeCallsImageModelDirectly(t *testing.T) { } } +func TestDesignAgentImageGeneratorModeBatchesSamePromptCount(t *testing.T) { + runner := &recordingRunner{} + creative := &recordingCreativeAgent{} + imageGenerator := &fakeImageGenerator{} + agent := NewDesignAgent(runner, imageGenerator, creative) + var emitted []design.Node + + plan, err := agent.PlanIncremental(context.Background(), design.AgentRequest{ + Project: design.Project{ + ID: "project-image-generator", + Nodes: []design.Node{ + { + ID: "generator-1", + Type: design.NodeTypeImage, + Tone: "image-generator", + Status: "generating", + LayerRole: "image-generator", + Width: 512, + Height: 512, + }, + }, + }, + Prompt: "画一只蓝色小鸟", + Mode: "image-generator", + ImageModel: "gpt-image-2", + ImageSize: "1024x1024", + TaskPlan: design.AgentTaskPlan{ + Intent: "image", + ImageCount: 2, + ImageTasks: []design.AgentImageTask{ + {Title: "蓝色小鸟 方案 1", Brief: "画一只蓝色小鸟"}, + {Title: "蓝色小鸟 方案 2", Brief: "画一只蓝色小鸟"}, + }, + }, + }, func(result design.AgentImageGenerationResult) error { + emitted = append(emitted, result.Node) + return nil + }) + if err != nil { + t.Fatal(err) + } + if runner.calls != 0 || creative.buildCalls != 0 { + t.Fatalf("expected direct image generator mode, runner=%d build=%d", runner.calls, creative.buildCalls) + } + if imageGenerator.callCount() != 1 { + t.Fatalf("expected one batched image request, got %d", imageGenerator.callCount()) + } + if imageGenerator.opts.Count != 2 { + t.Fatalf("expected image request n/count 2, got %#v", imageGenerator.opts) + } + if len(emitted) != 2 || len(plan.GeneratedNodes) != 2 { + t.Fatalf("expected two generated nodes, emitted=%#v plan=%#v", emitted, plan.GeneratedNodes) + } + if plan.GeneratedNodes[0].ID != "generator-1" || plan.GeneratedNodes[1].ID == "generator-1" { + t.Fatalf("expected first node to replace generator and second to be separate, got %#v", plan.GeneratedNodes) + } +} + func (g *fakeImageGenerator) callCount() int { g.mu.Lock() defer g.mu.Unlock() @@ -635,6 +693,7 @@ func TestGeneratedImageNodeFromImageGeneratorSlotBecomesRegularImage(t *testing. Tone: "image-generator", Status: "generating", LayerRole: "image-generator", + ParentID: "generator-parent", Width: 320, Height: 240, }, @@ -647,6 +706,9 @@ func TestGeneratedImageNodeFromImageGeneratorSlotBecomesRegularImage(t *testing. if node.LayerRole != "" || node.Tone != "visual" || node.Status != "success" { t.Fatalf("expected regular completed image node, got %#v", node) } + if node.ParentID != "" { + t.Fatalf("expected completed image generator node to clear parent id, got %#v", node) + } } func TestDesignAgentDoesNotRunFallbackAfterImageTimeout(t *testing.T) { diff --git a/server/internal/infrastructure/piagent/image_generation.go b/server/internal/infrastructure/piagent/image_generation.go index c826fd1..2dd2d49 100644 --- a/server/internal/infrastructure/piagent/image_generation.go +++ b/server/internal/infrastructure/piagent/image_generation.go @@ -136,7 +136,7 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq idempotencyKey := strings.TrimSpace(opts.IdempotencyKey) streamPlans := []bool{false} - if len(images) == 0 { + if len(images) == 0 && count == 1 { streamPlans = []bool{true, false} } diff --git a/server/internal/infrastructure/piagent/image_generation_test.go b/server/internal/infrastructure/piagent/image_generation_test.go index 8c38979..082a948 100644 --- a/server/internal/infrastructure/piagent/image_generation_test.go +++ b/server/internal/infrastructure/piagent/image_generation_test.go @@ -143,6 +143,33 @@ func TestImageClientSendsExplicitSize(t *testing.T) { } } +func TestImageClientSendsImageCount(t *testing.T) { + var requestBody map[string]any + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + 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-1.png"},{"url":"https://example.com/image-2.png"}]}`)) + })) + defer server.Close() + + client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"}) + image, err := client.Generate(context.Background(), "create two images", ImageRequestOptions{Model: "gpt-image-2", Count: 2}) + if err != nil { + t.Fatal(err) + } + if requestBody["n"] != float64(2) { + t.Fatalf("expected image request n=2, got %#v", requestBody["n"]) + } + if requestBody["stream"] != nil { + t.Fatalf("expected multi-image request to avoid stream, got %#v", requestBody["stream"]) + } + if len(image.URLs) != 2 { + t.Fatalf("expected two generated images, got %#v", image.URLs) + } +} + func TestImageClientUploadsDataReferenceImagesAsMultipart(t *testing.T) { var requestPath string var requestContentType string