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:
@@ -547,11 +547,12 @@ func (s *DesignService) AgentChat(ctx context.Context, projectID string, req des
|
||||
actionID := newID()
|
||||
targetNodeID := imageGeneratorTargetNodeID(req.Messages)
|
||||
isolatedImageGenerator := isImageGeneratorRequest(mode, targetNodeID)
|
||||
directImageMode := isDirectImageGenerationMode(mode, isolatedImageGenerator)
|
||||
enableWebSearch := req.EnableWebSearch || toolConfigEnablesSearch(req.ToolConfig)
|
||||
memory := buildAgentMemory(projectScopedToAgentThread(project, threadID), prompt, mode, req.Messages)
|
||||
if !isolatedImageGenerator {
|
||||
project.Messages = append(project.Messages,
|
||||
design.Message{
|
||||
nextMessages := []design.Message{
|
||||
{
|
||||
ID: newID(),
|
||||
Role: "user",
|
||||
Type: "user",
|
||||
@@ -562,8 +563,11 @@ func (s *DesignService) AgentChat(ctx context.Context, projectID string, req des
|
||||
StepID: newID(),
|
||||
CreatedAt: now,
|
||||
},
|
||||
agentPlanningMessage(threadID, actionID, prompt, enableWebSearch, now.Add(time.Millisecond)),
|
||||
)
|
||||
}
|
||||
if !directImageMode {
|
||||
nextMessages = append(nextMessages, agentPlanningMessage(threadID, actionID, prompt, enableWebSearch, now.Add(time.Millisecond)))
|
||||
}
|
||||
project.Messages = append(project.Messages, nextMessages...)
|
||||
project.Status = design.StatusExploring
|
||||
project.LastThreadID = threadID
|
||||
project.UpdatedAt = now
|
||||
@@ -571,31 +575,38 @@ func (s *DesignService) AgentChat(ctx context.Context, projectID string, req des
|
||||
if err := s.repo.Save(ctx, project); err != nil {
|
||||
return design.AgentThread{}, err
|
||||
}
|
||||
s.streamAgentPlanningMessage(ctx, project.ID, threadID, design.AgentTaskPlanRequest{
|
||||
Project: projectScopedToAgentThread(projectWithoutGeneratingNodes(project), threadID),
|
||||
Prompt: prompt,
|
||||
Mode: mode,
|
||||
Messages: req.Messages,
|
||||
Memory: buildAgentMemory(projectScopedToAgentThread(project, threadID), prompt, mode, req.Messages),
|
||||
WebSearchEnabled: enableWebSearch,
|
||||
})
|
||||
project, err = s.repo.Get(ctx, project.ID)
|
||||
if err != nil {
|
||||
return design.AgentThread{}, err
|
||||
}
|
||||
memory = buildAgentMemory(projectScopedToAgentThread(project, threadID), prompt, mode, req.Messages)
|
||||
if cancelled, err := s.projectThreadCancelled(ctx, project.ID, threadID); err != nil {
|
||||
return design.AgentThread{}, err
|
||||
} else if cancelled {
|
||||
return design.AgentThread{
|
||||
ThreadID: threadID,
|
||||
Status: "cancelled",
|
||||
ThreadIDType: threadIDType,
|
||||
Project: project,
|
||||
}, nil
|
||||
if !directImageMode {
|
||||
s.streamAgentPlanningMessage(ctx, project.ID, threadID, design.AgentTaskPlanRequest{
|
||||
Project: projectScopedToAgentThread(projectWithoutGeneratingNodes(project), threadID),
|
||||
Prompt: prompt,
|
||||
Mode: mode,
|
||||
Messages: req.Messages,
|
||||
Memory: buildAgentMemory(projectScopedToAgentThread(project, threadID), prompt, mode, req.Messages),
|
||||
WebSearchEnabled: enableWebSearch,
|
||||
})
|
||||
project, err = s.repo.Get(ctx, project.ID)
|
||||
if err != nil {
|
||||
return design.AgentThread{}, err
|
||||
}
|
||||
memory = buildAgentMemory(projectScopedToAgentThread(project, threadID), prompt, mode, req.Messages)
|
||||
if cancelled, err := s.projectThreadCancelled(ctx, project.ID, threadID); err != nil {
|
||||
return design.AgentThread{}, err
|
||||
} else if cancelled {
|
||||
return design.AgentThread{
|
||||
ThreadID: threadID,
|
||||
Status: "cancelled",
|
||||
ThreadIDType: threadIDType,
|
||||
Project: project,
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
taskPlan := s.planAgentTask(ctx, projectScopedToAgentThread(projectWithoutGeneratingNodes(project), threadID), prompt, mode, req.Messages, enableWebSearch, isolatedImageGenerator, memory)
|
||||
var taskPlan design.AgentTaskPlan
|
||||
if directImageMode {
|
||||
taskPlan = directImageTaskPlan(prompt, req.Messages)
|
||||
} else {
|
||||
taskPlan = s.planAgentTask(ctx, projectScopedToAgentThread(projectWithoutGeneratingNodes(project), threadID), prompt, mode, req.Messages, enableWebSearch, isolatedImageGenerator, memory)
|
||||
}
|
||||
if cancelled, err := s.projectThreadCancelled(ctx, project.ID, threadID); err != nil {
|
||||
return design.AgentThread{}, err
|
||||
} else if cancelled {
|
||||
@@ -606,7 +617,7 @@ func (s *DesignService) AgentChat(ctx context.Context, projectID string, req des
|
||||
Project: project,
|
||||
}, nil
|
||||
}
|
||||
if !isolatedImageGenerator {
|
||||
if !isolatedImageGenerator && !directImageMode {
|
||||
if err := s.maybeAddPlannedResearchStep(ctx, project.ID, threadID, projectScopedToAgentThread(project, threadID), prompt, mode, req.Messages, enableWebSearch, taskPlan); err != nil {
|
||||
return design.AgentThread{}, err
|
||||
}
|
||||
@@ -647,7 +658,9 @@ func (s *DesignService) AgentChat(ctx context.Context, projectID string, req des
|
||||
assistantMessageID := ""
|
||||
if !isolatedImageGenerator {
|
||||
if shouldGenerateImage {
|
||||
applyPlannerTaskPlanMessage(&project, threadID, prompt, taskPlan)
|
||||
if !directImageMode {
|
||||
applyPlannerTaskPlanMessage(&project, threadID, prompt, taskPlan)
|
||||
}
|
||||
} else {
|
||||
assistantMessageID = preparePlannerForChatMessage(&project, threadID, actionID, prompt, now.Add(time.Millisecond))
|
||||
}
|
||||
@@ -3468,6 +3481,84 @@ func shouldGenerateForAgentPrompt(prompt string, mode string, messages []design.
|
||||
return false
|
||||
}
|
||||
|
||||
func isDirectImageGenerationMode(mode string, isolated bool) bool {
|
||||
if isolated {
|
||||
return true
|
||||
}
|
||||
modeText := strings.ToLower(strings.TrimSpace(mode))
|
||||
return modeText == "image" || modeText == "pure-image" || modeText == "image-generator" || strings.Contains(modeText, "image-action")
|
||||
}
|
||||
|
||||
func directImageTaskPlan(prompt string, messages []design.AgentChatMessage) design.AgentTaskPlan {
|
||||
count := imageModeSettingsCount(messages)
|
||||
if count <= 0 {
|
||||
count = 1
|
||||
}
|
||||
if count > 10 {
|
||||
count = 10
|
||||
}
|
||||
tasks := make([]design.AgentImageTask, 0, count)
|
||||
for index := 0; index < count; index++ {
|
||||
tasks = append(tasks, design.AgentImageTask{
|
||||
Title: generatedTaskFallbackTitle(prompt, index, count),
|
||||
Brief: prompt,
|
||||
})
|
||||
}
|
||||
return design.AgentTaskPlan{
|
||||
Intent: "image",
|
||||
Reason: "direct image mode",
|
||||
ImageCount: count,
|
||||
ImageTasks: tasks,
|
||||
}
|
||||
}
|
||||
|
||||
func imageModeSettingsCount(messages []design.AgentChatMessage) int {
|
||||
for _, message := range messages {
|
||||
for _, content := range message.Contents {
|
||||
if !strings.EqualFold(strings.TrimSpace(content.TextSource), "settings") {
|
||||
continue
|
||||
}
|
||||
if count := imageModeSettingsCountFromText(content.Text); count > 0 {
|
||||
return count
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func imageModeSettingsCountFromText(text string) int {
|
||||
lower := strings.ToLower(text)
|
||||
for _, marker := range []string{"数量", "count"} {
|
||||
index := strings.Index(lower, marker)
|
||||
if index < 0 {
|
||||
continue
|
||||
}
|
||||
if count := firstPositiveInt(lower[index+len(marker):]); count > 0 {
|
||||
return count
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func firstPositiveInt(text string) int {
|
||||
value := 0
|
||||
seen := false
|
||||
for _, char := range text {
|
||||
if char < '0' || char > '9' {
|
||||
if seen {
|
||||
break
|
||||
}
|
||||
continue
|
||||
}
|
||||
seen = true
|
||||
value = value*10 + int(char-'0')
|
||||
if value > 10 {
|
||||
return 10
|
||||
}
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func shouldForceImageMode(mode string) bool {
|
||||
modeText := strings.ToLower(strings.TrimSpace(mode))
|
||||
return strings.Contains(modeText, "image-action") || modeText == "image" || modeText == "pure-image"
|
||||
@@ -4048,9 +4139,6 @@ func (s *DesignService) clearStaleGeneratingPlaceholders(project design.Project,
|
||||
|
||||
func (s *DesignService) restoreFinishedGeneratedImageNodes(project design.Project, now time.Time) (design.Project, bool) {
|
||||
artifacts := generatedImageArtifacts(project.Messages)
|
||||
if len(artifacts) == 0 {
|
||||
return project, false
|
||||
}
|
||||
artifactsByID := generatedImageArtifactsByID(project.Messages)
|
||||
usedArtifacts := make(map[string]bool)
|
||||
next := make([]design.Node, 0, len(project.Nodes))
|
||||
@@ -4068,6 +4156,10 @@ func (s *DesignService) restoreFinishedGeneratedImageNodes(project design.Projec
|
||||
}
|
||||
changed = true
|
||||
}
|
||||
if normalized, ok := completedImageGeneratorNode(node); ok {
|
||||
node = normalized
|
||||
changed = true
|
||||
}
|
||||
if project.Status == design.StatusReady && node.Type == design.NodeTypeImage && node.Status == "generating" && !isGeneratedImageContent(node.Content) {
|
||||
changed = true
|
||||
continue
|
||||
@@ -4155,6 +4247,17 @@ func restoreNodeFromGeneratedArtifact(node design.Node, artifact design.Generato
|
||||
return node
|
||||
}
|
||||
|
||||
func completedImageGeneratorNode(node design.Node) (design.Node, bool) {
|
||||
if node.Type != design.NodeTypeImage || node.LayerRole != "image-generator" || node.Status != "success" || !isGeneratedImageContent(node.Content) {
|
||||
return node, false
|
||||
}
|
||||
node.LayerRole = ""
|
||||
if strings.TrimSpace(node.Tone) == "" || node.Tone == "image-generator" {
|
||||
node.Tone = "visual"
|
||||
}
|
||||
return node, true
|
||||
}
|
||||
|
||||
func generationPlaceholdersAreStale(project design.Project, now time.Time) bool {
|
||||
if !hasGeneratingCanvasWork(project.Nodes) {
|
||||
return false
|
||||
@@ -4361,20 +4464,24 @@ func imageGeneratorTargetNodeID(messages []design.AgentChatMessage) string {
|
||||
|
||||
func mergeGeneratedNodesIntoTarget(nodes []design.Node, generated []design.Node, targetNodeID string) []design.Node {
|
||||
if strings.TrimSpace(targetNodeID) == "" || len(generated) == 0 {
|
||||
return append(nodes, generated...)
|
||||
return append(nodes, completedImageGeneratorNodes(generated)...)
|
||||
}
|
||||
targetNodeID = strings.TrimSpace(targetNodeID)
|
||||
targetIndex := -1
|
||||
for index, node := range nodes {
|
||||
if node.ID == targetNodeID && node.Type == design.NodeTypeImage && node.LayerRole == "image-generator" {
|
||||
if node.ID == targetNodeID && node.Type == design.NodeTypeImage {
|
||||
targetIndex = index
|
||||
break
|
||||
}
|
||||
}
|
||||
if targetIndex < 0 {
|
||||
return append(nodes, generated...)
|
||||
return append(nodes, completedImageGeneratorNodes(generated)...)
|
||||
}
|
||||
|
||||
primary := generated[0]
|
||||
if normalized, ok := completedImageGeneratorNode(primary); ok {
|
||||
primary = normalized
|
||||
}
|
||||
target := nodes[targetIndex]
|
||||
centerX := target.X + target.Width/2
|
||||
centerY := target.Y + target.Height/2
|
||||
@@ -4383,6 +4490,9 @@ func mergeGeneratedNodesIntoTarget(nodes []design.Node, generated []design.Node,
|
||||
target.Tone = primary.Tone
|
||||
target.Status = primary.Status
|
||||
target.LayerRole = ""
|
||||
if strings.TrimSpace(target.Tone) == "" || target.Tone == "image-generator" {
|
||||
target.Tone = "visual"
|
||||
}
|
||||
if primary.Width > 0 {
|
||||
target.Width = primary.Width
|
||||
target.X = centerX - target.Width/2
|
||||
@@ -4393,11 +4503,23 @@ func mergeGeneratedNodesIntoTarget(nodes []design.Node, generated []design.Node,
|
||||
}
|
||||
nodes[targetIndex] = target
|
||||
if len(generated) > 1 {
|
||||
return append(nodes, generated[1:]...)
|
||||
return append(nodes, completedImageGeneratorNodes(generated[1:])...)
|
||||
}
|
||||
return nodes
|
||||
}
|
||||
|
||||
func completedImageGeneratorNodes(nodes []design.Node) []design.Node {
|
||||
next := make([]design.Node, len(nodes))
|
||||
for index, node := range nodes {
|
||||
if normalized, ok := completedImageGeneratorNode(node); ok {
|
||||
next[index] = normalized
|
||||
continue
|
||||
}
|
||||
next[index] = node
|
||||
}
|
||||
return next
|
||||
}
|
||||
|
||||
func findNodeIndex(nodes []design.Node, nodeID string) int {
|
||||
for index, node := range nodes {
|
||||
if node.ID == nodeID {
|
||||
|
||||
@@ -554,6 +554,75 @@ func TestAgentChatRequiresTextPrompt(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentChatImageModeBypassesPlannerAndHidesSettings(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
repo := memory.NewProjectRepository()
|
||||
planner := &recordingTaskPlanner{
|
||||
fakeCreativeChatAgent: fakeCreativeChatAgent{
|
||||
decision: design.AgentDecision{Intent: "image"},
|
||||
},
|
||||
}
|
||||
service := NewDesignService(repo, nil, nil, nil, nil, planner)
|
||||
queue := &fakeJobQueue{}
|
||||
service.SetJobQueue(queue)
|
||||
project := design.Project{
|
||||
ID: "project-direct-image",
|
||||
Title: "Direct image",
|
||||
Status: design.StatusReady,
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
if err := repo.Save(ctx, project); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
thread, err := service.AgentChat(ctx, project.ID, design.AgentChatRequest{
|
||||
Mode: "image",
|
||||
ImageSize: "1536x1024",
|
||||
Messages: []design.AgentChatMessage{
|
||||
{
|
||||
Role: "user",
|
||||
Contents: []design.AgentContent{
|
||||
{Type: "text", Text: "a cute bird flatten ui", TextSource: "input"},
|
||||
{Type: "text", Text: "图片生成参数:质量 自动,尺寸 1536x1024,比例 3:2,数量 3 张。纯生图模式,请直接围绕用户输入生成图片,不展开额外聊天。", TextSource: "settings"},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if thread.Status != "running" {
|
||||
t.Fatalf("expected running thread, got %s", thread.Status)
|
||||
}
|
||||
if planner.called != 0 {
|
||||
t.Fatalf("expected image mode to bypass planner, got %d calls", planner.called)
|
||||
}
|
||||
if len(queue.jobs) != 1 {
|
||||
t.Fatalf("expected one image job, got %#v", queue.jobs)
|
||||
}
|
||||
job := queue.jobs[0]
|
||||
if job.TaskPlan.ImageCount != 3 || len(job.TaskPlan.ImageTasks) != 3 {
|
||||
t.Fatalf("expected direct image count from settings, got %#v", job.TaskPlan)
|
||||
}
|
||||
if job.ImageSize != "1536x1024" {
|
||||
t.Fatalf("expected image size to be preserved, got %s", job.ImageSize)
|
||||
}
|
||||
|
||||
saved, err := repo.Get(ctx, project.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(saved.Messages) != 1 {
|
||||
t.Fatalf("expected only the user message before generation starts, got %#v", saved.Messages)
|
||||
}
|
||||
if saved.Messages[0].Content != "a cute bird flatten ui" {
|
||||
t.Fatalf("expected settings to stay out of visible messages, got %q", saved.Messages[0].Content)
|
||||
}
|
||||
if len(saved.Nodes) != 3 {
|
||||
t.Fatalf("expected three direct image placeholders, got %#v", saved.Nodes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentChatCreatesPlaceholdersFromModelPlannedImageTasks(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
repo := memory.NewProjectRepository()
|
||||
|
||||
@@ -50,6 +50,44 @@ func TestMergeGeneratedNodesIntoImageGeneratorTargetUsesGeneratedSizeAndKeepsCen
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeGeneratedNodesIntoCompletedImageGeneratorTargetDoesNotReappend(t *testing.T) {
|
||||
nodes := []design.Node{
|
||||
{
|
||||
ID: "generator-1",
|
||||
Type: design.NodeTypeImage,
|
||||
Title: "旧图",
|
||||
Content: "https://example.com/old.png",
|
||||
Tone: "visual",
|
||||
Status: "success",
|
||||
LayerRole: "",
|
||||
Width: 360,
|
||||
Height: 240,
|
||||
},
|
||||
}
|
||||
generated := []design.Node{
|
||||
{
|
||||
ID: "generator-1",
|
||||
Type: design.NodeTypeImage,
|
||||
Title: "新图",
|
||||
Content: "https://example.com/new.png",
|
||||
Tone: "image-generator",
|
||||
Status: "success",
|
||||
LayerRole: "image-generator",
|
||||
Width: 1024,
|
||||
Height: 1024,
|
||||
},
|
||||
}
|
||||
|
||||
merged := mergeGeneratedNodesIntoTarget(nodes, generated, "generator-1")
|
||||
if len(merged) != 1 {
|
||||
t.Fatalf("expected existing target to be updated in place, got %#v", merged)
|
||||
}
|
||||
node := merged[0]
|
||||
if node.Content != "https://example.com/new.png" || node.LayerRole != "" || node.Tone != "visual" {
|
||||
t.Fatalf("expected completed image generator to become regular image, got %#v", node)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageGeneratorTargetNodeIDParsesDirectiveCaseInsensitively(t *testing.T) {
|
||||
targetID := imageGeneratorTargetNodeID([]design.AgentChatMessage{
|
||||
{
|
||||
@@ -64,6 +102,36 @@ func TestImageGeneratorTargetNodeIDParsesDirectiveCaseInsensitively(t *testing.T
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreFinishedGeneratedImageNodesConvertsCompletedImageGeneratorNode(t *testing.T) {
|
||||
now := time.Now()
|
||||
project := design.Project{
|
||||
ID: "project-completed-generator",
|
||||
Status: design.StatusReady,
|
||||
Nodes: []design.Node{
|
||||
{
|
||||
ID: "generator-1",
|
||||
Type: design.NodeTypeImage,
|
||||
Title: "生成图片",
|
||||
Content: "http://localhost:19000/canvas-assets/cat.png",
|
||||
Tone: "image-generator",
|
||||
Status: "success",
|
||||
LayerRole: "image-generator",
|
||||
Width: 1024,
|
||||
Height: 1024,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
restored, changed := (*DesignService)(nil).restoreFinishedGeneratedImageNodes(project, now.Add(time.Second))
|
||||
if !changed {
|
||||
t.Fatal("expected completed image generator node to be normalized")
|
||||
}
|
||||
node := restored.Nodes[0]
|
||||
if node.LayerRole != "" || node.Tone != "visual" || node.Status != "success" || node.Content != "http://localhost:19000/canvas-assets/cat.png" {
|
||||
t.Fatalf("expected regular completed image node, got %#v", node)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreFinishedGeneratedImageNodesRepairsStaleCanvasSnapshot(t *testing.T) {
|
||||
now := time.Now()
|
||||
project := design.Project{
|
||||
|
||||
@@ -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