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"