root/moteva
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(piagent): batch-generate images in one call when generators share a prompt

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
liangxu
2026-07-09 09:05:18 +08:00
parent 8a28a0bcea
commit cc30ae4284
4 changed files with 223 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/infrastructure/piagent/design_agent.go
+133 -4
View File
@@ -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"
server/internal/infrastructure/piagent/design_agent_test.go
+62
View File
@@ -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) {
server/internal/infrastructure/piagent/image_generation.go
+1 -1
View File
@@ -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}
}
server/internal/infrastructure/piagent/image_generation_test.go
+27
View File
@@ -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