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:
@@ -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{})
|
||||
|
||||
Reference in New Issue
Block a user