feat(piagent): bypass planner for direct image generation modes

For image / pure-image / image-generator / image-action modes, skip the
task-planning round trip and hand the prompt straight to the image model,
deriving the requested image count from the composer settings. Also
normalize completed image-generator nodes back into regular image nodes
(clearing the image-generator layer role and defaulting Tone to "visual")
when merging generated results and when restoring finished nodes, so a
finished generator node is not re-appended as a duplicate.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-09 01:09:59 +08:00
parent a39e18950c
commit ffd8d78df7
5 changed files with 573 additions and 60 deletions
@@ -46,6 +46,10 @@ func NewDesignAgent(base design.AgentRunner, imageAgent ImageGenerator, textAgen
}
func (a *DesignAgent) Plan(ctx context.Context, req design.AgentRequest) (design.AgentPlan, error) {
if shouldUseDirectImageGeneration(req) {
return a.planDirectImageGeneration(ctx, req)
}
plan, err := a.base.Plan(ctx, req)
if err != nil {
return design.AgentPlan{}, err
@@ -76,6 +80,10 @@ func (a *DesignAgent) Plan(ctx context.Context, req design.AgentRequest) (design
}
func (a *DesignAgent) PlanIncremental(ctx context.Context, req design.AgentRequest, onResult func(design.AgentImageGenerationResult) error) (design.AgentPlan, error) {
if shouldUseDirectImageGeneration(req) {
return a.planDirectImageGenerationIncremental(ctx, req, onResult)
}
plan, err := a.base.Plan(ctx, req)
if err != nil {
return design.AgentPlan{}, err
@@ -115,6 +123,92 @@ func (a *DesignAgent) PlanIncremental(ctx context.Context, req design.AgentReque
return plan, nil
}
func (a *DesignAgent) planDirectImageGeneration(ctx context.Context, req design.AgentRequest) (design.AgentPlan, error) {
if a.imageAgent == nil {
return design.AgentPlan{}, fmt.Errorf("gpt image 2 is not configured")
}
generationReq := req
generationReq.ImageSize = resolveImageSizeForGeneration(req.ImageSize, req.Prompt)
prompts := directImagePrompts(generationReq)
imageURLs, err := a.generatePlannedImages(ctx, prompts, generationReq)
if err != nil {
return design.AgentPlan{}, fmt.Errorf("gpt image 2 generation failed: %w", err)
}
plan := directImagePlan(generationReq)
applyPlannedGeneratedImages(&plan, generationReq, imageURLs, prompts)
title := ""
if len(prompts) > 0 {
title = prompts[0].Title
}
plan.Message.Title = generatedMessageTitle(title, req.Prompt)
return plan, nil
}
func (a *DesignAgent) planDirectImageGenerationIncremental(ctx context.Context, req design.AgentRequest, onResult func(design.AgentImageGenerationResult) error) (design.AgentPlan, error) {
if a.imageAgent == nil {
return design.AgentPlan{}, fmt.Errorf("gpt image 2 is not configured")
}
generationReq := req
generationReq.ImageSize = resolveImageSizeForGeneration(req.ImageSize, req.Prompt)
prompts := directImagePrompts(generationReq)
plan := directImagePlan(generationReq)
nodes, err := a.generatePlannedImagesIncremental(ctx, prompts, generationReq, 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
title := ""
if len(prompts) > 0 {
title = prompts[0].Title
}
plan.Message.Title = generatedMessageTitle(title, req.Prompt)
return plan, nil
}
func directImagePlan(req design.AgentRequest) design.AgentPlan {
return design.AgentPlan{
Message: design.Message{
ID: uuid.NewString(),
Role: "assistant",
Type: "assistant",
Title: generatedMessageTitle("", req.Prompt),
Content: "",
},
}
}
func directImagePrompts(req design.AgentRequest) []design.AgentImagePrompt {
count := plannedGenerationCount(req)
prompt := strings.TrimSpace(req.Prompt)
if prompt == "" {
prompt = fallbackDirectImagePrompt(req)
}
prompts := make([]design.AgentImagePrompt, 0, count)
for index := 0; index < count; index++ {
title := normalizeGeneratedTitle("", req.Prompt)
if index < len(req.TaskPlan.ImageTasks) {
if taskTitle := strings.TrimSpace(req.TaskPlan.ImageTasks[index].Title); taskTitle != "" {
title = normalizeGeneratedTitle(taskTitle, req.Prompt)
}
}
prompts = append(prompts, design.AgentImagePrompt{
Title: title,
Prompt: prepareGPTImagePrompt(variantImagePrompt(prompt, index+1, count), imageGenerationPromptMaxRunes),
})
}
return prompts
}
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")
}
func (a *DesignAgent) buildImagePrompts(ctx context.Context, req design.AgentRequest) []design.AgentImagePrompt {
count := plannedGenerationCount(req)
prompts := make([]design.AgentImagePrompt, 0, count)
@@ -667,28 +761,7 @@ func applyGeneratedImages(plan *design.AgentPlan, req design.AgentRequest, image
slots := generatedImageSlots(req, plan.GeneratedNodes, requestedCount)
nodes := make([]design.Node, 0, requestedCount)
for index := 0; index < requestedCount; index++ {
node := slots[index]
node.Type = design.NodeTypeImage
node.Title = generatedNodeTitle(title, req.Prompt, index, requestedCount)
node.Content = imageURLs[index]
node.Status = "success"
if node.ID == "" {
node.ID = uuid.NewString()
}
if width, height, ok := imageSizeDimensions(req.ImageSize); ok {
node.Width = width
node.Height = height
}
if node.Width <= 0 {
node.Width = design.DefaultGeneratedImageWidth
}
if node.Height <= 0 {
node.Height = design.DefaultGeneratedImageHeight
}
if node.Tone == "" {
node.Tone = "visual"
}
nodes = append(nodes, node)
nodes = append(nodes, generatedImageNodeFromSlot(slots[index], req, imageURLs[index], title, index, requestedCount))
}
plan.GeneratedNodes = nodes
plan.Connections = nil
@@ -731,7 +804,10 @@ func generatedImageNodeFromSlot(slot design.Node, req design.AgentRequest, image
if node.Height <= 0 {
node.Height = design.DefaultGeneratedImageHeight
}
if node.Tone == "" {
if node.LayerRole == "image-generator" {
node.LayerRole = ""
}
if node.Tone == "" || node.Tone == "image-generator" {
node.Tone = "visual"
}
return node
@@ -50,6 +50,42 @@ func (fakeRunner) Replay(ctx context.Context, project design.Project) (design.Ag
return design.AgentReplay{ProjectID: project.ID, Messages: project.Messages}, nil
}
type recordingRunner struct {
calls int
}
func (r *recordingRunner) Plan(ctx context.Context, req design.AgentRequest) (design.AgentPlan, error) {
r.calls++
return design.AgentPlan{Message: design.Message{ID: "message-1", Role: "assistant", Title: "done", Content: "done"}}, nil
}
func (r *recordingRunner) Replay(ctx context.Context, project design.Project) (design.AgentReplay, error) {
return design.AgentReplay{ProjectID: project.ID, Messages: project.Messages}, nil
}
type recordingCreativeAgent struct {
buildCalls int
summarizeCalls int
}
func (a *recordingCreativeAgent) Decide(ctx context.Context, req design.AgentDecisionRequest) (design.AgentDecision, error) {
return design.AgentDecision{}, nil
}
func (a *recordingCreativeAgent) Respond(ctx context.Context, req design.AgentConversationRequest) (string, error) {
return "", nil
}
func (a *recordingCreativeAgent) SummarizeImage(ctx context.Context, req design.AgentImageSummaryRequest) (string, error) {
a.summarizeCalls++
return "summary", nil
}
func (a *recordingCreativeAgent) BuildImagePrompt(ctx context.Context, req design.AgentImagePromptRequest) (design.AgentImagePrompt, error) {
a.buildCalls++
return design.AgentImagePrompt{Title: "Rewritten", Prompt: "rewritten by gpt-5.5"}, nil
}
type fakeImageGenerator struct {
mu sync.Mutex
called bool
@@ -114,6 +150,126 @@ func imageOptionIndex(prompt string, fallback int) int {
return fallback
}
func TestDesignAgentImageModeCallsImageModelDirectly(t *testing.T) {
runner := &recordingRunner{}
creative := &recordingCreativeAgent{}
imageGenerator := &fakeImageGenerator{}
agent := NewDesignAgent(runner, imageGenerator, creative)
plan, err := agent.Plan(context.Background(), design.AgentRequest{
Project: design.Project{ID: "project-direct", Title: "Direct image"},
Prompt: "a cute bird flatten ui",
Mode: "image",
ImageModel: "gpt-image-2",
ImageSize: "1536x1024",
TaskPlan: design.AgentTaskPlan{
Intent: "image",
ImageCount: 1,
ImageTasks: []design.AgentImageTask{
{Title: "Cute bird", Brief: "a cute bird flatten ui"},
},
},
})
if err != nil {
t.Fatal(err)
}
if runner.calls != 0 {
t.Fatalf("expected direct image mode to bypass base runner, got %d calls", runner.calls)
}
if creative.buildCalls != 0 || creative.summarizeCalls != 0 {
t.Fatalf("expected direct image mode to bypass gpt-5.5 helpers, got build=%d summarize=%d", creative.buildCalls, creative.summarizeCalls)
}
if !imageGenerator.called {
t.Fatal("expected image generator to be called")
}
if !strings.Contains(imageGenerator.prompt, "a cute bird flatten ui") {
t.Fatalf("expected raw user prompt to reach image generator, got %q", imageGenerator.prompt)
}
if strings.Contains(imageGenerator.prompt, "rewritten by gpt-5.5") || strings.Contains(imageGenerator.prompt, "Create one polished") {
t.Fatalf("expected direct prompt without analysis wrapper, got %q", imageGenerator.prompt)
}
if len(plan.GeneratedNodes) != 1 || plan.GeneratedNodes[0].Content == "" {
t.Fatalf("expected one generated image node, got %#v", plan.GeneratedNodes)
}
}
func TestDesignAgentImageGeneratorModeCallsImageModelDirectly(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",
Messages: []design.AgentChatMessage{
{
Role: "user",
Contents: []design.AgentContent{
{Type: "image", ImageURL: "https://example.com/reference.png"},
{Type: "text", Text: "画一只蓝色小鸟", TextSource: "input"},
{Type: "text", Text: "Image generator target node: generator-1", TextSource: "settings"},
},
},
},
TaskPlan: design.AgentTaskPlan{
Intent: "image",
ImageCount: 1,
ImageTasks: []design.AgentImageTask{
{Title: "蓝色小鸟", Brief: "画一只蓝色小鸟"},
},
},
}, func(result design.AgentImageGenerationResult) error {
emitted = append(emitted, result.Node)
return nil
})
if err != nil {
t.Fatal(err)
}
if runner.calls != 0 {
t.Fatalf("expected image generator mode to bypass base runner, got %d calls", runner.calls)
}
if creative.buildCalls != 0 || creative.summarizeCalls != 0 {
t.Fatalf("expected image generator mode to bypass gpt-5.5 helpers, got build=%d summarize=%d", creative.buildCalls, creative.summarizeCalls)
}
if !imageGenerator.called {
t.Fatal("expected image generator to be called")
}
if !strings.Contains(imageGenerator.prompt, "画一只蓝色小鸟") {
t.Fatalf("expected raw image generator prompt to reach image generator, got %q", imageGenerator.prompt)
}
if strings.Contains(imageGenerator.prompt, "rewritten by gpt-5.5") || strings.Contains(imageGenerator.prompt, "Create one polished") {
t.Fatalf("expected direct image generator prompt without analysis wrapper, got %q", imageGenerator.prompt)
}
if len(imageGenerator.opts.Images) != 1 || imageGenerator.opts.Images[0] != "https://example.com/reference.png" {
t.Fatalf("expected reference image to be passed directly, got %#v", imageGenerator.opts.Images)
}
if len(emitted) != 1 || len(plan.GeneratedNodes) != 1 {
t.Fatalf("expected one generated node, emitted=%#v plan=%#v", emitted, plan.GeneratedNodes)
}
node := plan.GeneratedNodes[0]
if node.ID != "generator-1" || node.LayerRole != "" || node.Tone != "visual" || node.Status != "success" {
t.Fatalf("expected image generator slot to become regular image, got %#v", node)
}
}
func (g *fakeImageGenerator) callCount() int {
g.mu.Lock()
defer g.mu.Unlock()
@@ -168,7 +324,7 @@ func TestDesignAgentGeneratesImageThroughPiAgent(t *testing.T) {
plan, err := agent.Plan(context.Background(), design.AgentRequest{
Prompt: "生成一张 still day 官网首页图片",
Mode: "image",
Mode: "design",
ImageModel: "gpt-image-2",
ImageSize: "1536x1024",
})
@@ -471,6 +627,28 @@ func TestGeneratedImageSlotsIgnoreCanvasActionPlaceholders(t *testing.T) {
}
}
func TestGeneratedImageNodeFromImageGeneratorSlotBecomesRegularImage(t *testing.T) {
node := generatedImageNodeFromSlot(
design.Node{
ID: "generator-1",
Type: design.NodeTypeImage,
Tone: "image-generator",
Status: "generating",
LayerRole: "image-generator",
Width: 320,
Height: 240,
},
design.AgentRequest{Prompt: "生成一张小猫图片"},
"https://example.com/cat.png",
"小猫",
0,
1,
)
if node.LayerRole != "" || node.Tone != "visual" || node.Status != "success" {
t.Fatalf("expected regular completed image node, got %#v", node)
}
}
func TestDesignAgentDoesNotRunFallbackAfterImageTimeout(t *testing.T) {
imageGenerator := &failingImageGenerator{err: context.DeadlineExceeded}
agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{})