重构提取提示词生成逻辑，移除文本拼接方式，新增 YAML 序列化结构与测试

2026-09-11 14:54:00 +08:00
parent 1303592a86
commit 7c7d55dcab
3 changed files with 124 additions and 86 deletions
@@ -252,37 +252,6 @@ mod tests {
        assert!(next_image_batch(&separation, &separation.trees[0]).is_empty());
    }

-    #[test]
-    fn next_extract_prompt_contains_previous_rework_notes() {
-        let node = SeparationNode {
-            id: NodeId::new("image").unwrap(),
-            kind: SeparationNodeKind::ImageTarget,
-            global_pos_x_px: 0,
-            global_pos_y_px: 0,
-            width_px: 1,
-            height_px: 1,
-            note: SeparationNote {
-                description: "图标".to_string(),
-                rework_notes: vec!["不要带父背景".to_string()],
-            },
-            children: vec![],
-            rework_count: 0,
-        };
-        let tree = SeparationTree {
-            src_ui_design: UIDesignImageId::new("page").unwrap(),
-            root: node.clone(),
-            root_extractable: true,
-        };
-        let separation = SeparationState {
-            schema_version: SEPARATION_STATE_SCHEMA_VERSION.to_string(),
-            trees: vec![tree.clone()],
-            bound: vec![],
-            problematic_nodes: vec![],
-        };
-        let prompt = super::prompt::gen_extract_prompt(&separation, &tree, &[&node]);
-        assert!(prompt.contains("previous rework notes:\n- 不要带父背景"));
-    }
-
    #[test]
    fn binding_validation_rejects_overlong_rework_note() {
        let node = SeparationNode {
@@ -1,8 +1,45 @@
+use crate::ui_editor::commands::separation::model::{SeparationState, SeparationTree};
 use crate::ui_editor::commands::separation::{SeparationNode, SeparationNodeKind};
+use crate::ui_editor::utils::NodeId;
+use serde::Serialize;
 use std::collections::HashSet;

-use crate::ui_editor::commands::separation::model::{SeparationState, SeparationTree};
-use crate::ui_editor::utils::NodeId;
+#[derive(Serialize)]
+struct ExtractPromptDocument {
+    ui_layer_tree: ExtractPromptNode,
+}
+
+#[derive(Serialize)]
+struct ExtractPromptNode {
+    index: usize,
+    status: ExtractPromptStatus,
+    rect: ExtractPromptRect,
+    description: String,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    rework_notes: Vec<String>,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    children: Vec<ExtractPromptNode>,
+}
+
+#[derive(Serialize)]
+enum ExtractPromptStatus {
+    #[serde(rename = "OUTPUT_THIS_TURN")]
+    OutputThisTurn,
+    #[serde(rename = "DONE")]
+    Done,
+    #[serde(rename = "CONTEXT_ONLY")]
+    ContextOnly,
+    #[serde(rename = "REMOVE_ONLY")]
+    RemoveOnly,
+}
+
+#[derive(Serialize)]
+struct ExtractPromptRect {
+    x: u32,
+    y: u32,
+    width: u32,
+    height: u32,
+}

 pub(crate) fn gen_extract_prompt(
    state: &SeparationState,
@@ -36,66 +73,97 @@ pub(crate) fn gen_extract_prompt(
                .map(|node| node.node_id.clone()),
        )
        .collect::<HashSet<_>>();
-    let mut lines = Vec::new();
-    let mut index = 1usize;
-    append_context_lines(
-        &tree.root,
-        0,
-        true,
-        &target_ids,
-        &terminal_ids,
-        &mut index,
-        &mut lines,
-    );
-    result.push_str(&lines.join("\n"));
+    let mut index = 1;
+    let document = ExtractPromptDocument {
+        ui_layer_tree: project_node(&tree.root, &target_ids, &terminal_ids, &mut index),
+    };
+    let yaml = serde_yaml::to_string(&document)
+        .expect("UI separation extract prompt projection must be serializable");
+
+    result.push_str("```yaml\n");
+    result.push_str(&yaml);
+    result.push_str("```\n");
    result
 }

-fn append_context_lines(
+fn project_node(
    node: &SeparationNode,
-    depth: usize,
-    is_root: bool,
    target_ids: &HashSet<NodeId>,
    terminal_ids: &HashSet<NodeId>,
    index: &mut usize,
-    lines: &mut Vec<String>,
-) {
-    let status = match node.kind {
-        SeparationNodeKind::TextRemovalOnly => "REMOVE_ONLY",
-        SeparationNodeKind::PureContainer => "CONTEXT_ONLY",
-        SeparationNodeKind::ImageTarget => {
-            if target_ids.contains(&node.id) {
-                "OUTPUT_THIS_TURN"
-            } else if terminal_ids.contains(&node.id) {
-                "DONE"
-            } else {
-                "CONTEXT_ONLY"
-            }
-        }
-    };
-    let role = if is_root { "root" } else { "child" };
-    lines.push(format!(
-        "{}. [{}] depth={} role={} rect=({}, {}, {}, {}) {}",
-        *index,
-        status,
-        depth,
-        role,
-        node.global_pos_x_px,
-        node.global_pos_y_px,
-        node.width_px,
-        node.height_px,
-        node.note.as_prompt()
-    ));
+) -> ExtractPromptNode {
+    let current_index = *index;
    *index += 1;
-    for child in &node.children {
-        append_context_lines(
-            child,
-            depth + 1,
-            false,
-            target_ids,
-            terminal_ids,
-            index,
-            lines,
-        );
+
+    let status = match node.kind {
+        SeparationNodeKind::TextRemovalOnly => ExtractPromptStatus::RemoveOnly,
+        SeparationNodeKind::PureContainer => ExtractPromptStatus::ContextOnly,
+        SeparationNodeKind::ImageTarget if target_ids.contains(&node.id) => {
+            ExtractPromptStatus::OutputThisTurn
+        }
+        SeparationNodeKind::ImageTarget if terminal_ids.contains(&node.id) => {
+            ExtractPromptStatus::Done
+        }
+        SeparationNodeKind::ImageTarget => ExtractPromptStatus::ContextOnly,
+    };
+
+    let children = node
+        .children
+        .iter()
+        .map(|child| project_node(child, target_ids, terminal_ids, index))
+        .collect();
+
+    ExtractPromptNode {
+        index: current_index,
+        status,
+        rect: ExtractPromptRect {
+            x: node.global_pos_x_px,
+            y: node.global_pos_y_px,
+            width: node.width_px,
+            height: node.height_px,
+        },
+        description: node.note.description.clone(),
+        rework_notes: node.note.rework_notes.clone(),
+        children,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ui_editor::commands::separation::model::SeparationNote;
+
+    #[test]
+    fn projects_source_node_to_prompt_view() {
+        let node = SeparationNode {
+            id: NodeId::new("image").unwrap(),
+            kind: SeparationNodeKind::ImageTarget,
+            global_pos_x_px: 1,
+            global_pos_y_px: 2,
+            width_px: 3,
+            height_px: 4,
+            note: SeparationNote {
+                description: "按钮".to_string(),
+                rework_notes: vec!["保留圆角".to_string()],
+            },
+            children: vec![],
+            rework_count: 0,
+        };
+        let mut index = 1;
+        let target_ids = HashSet::from([node.id.clone()]);
+        let projected = project_node(&node, &target_ids, &HashSet::new(), &mut index);
+
+        assert_eq!(projected.index, 1);
+        assert!(matches!(
+            projected.status,
+            ExtractPromptStatus::OutputThisTurn
+        ));
+        assert_eq!(projected.rect.x, 1);
+        assert_eq!(projected.rect.y, 2);
+        assert_eq!(projected.rect.width, 3);
+        assert_eq!(projected.rect.height, 4);
+        assert_eq!(projected.description, "按钮");
+        assert_eq!(projected.rework_notes, vec!["保留圆角".to_string()]);
+        assert!(projected.children.is_empty());
    }
 }