Skip to main content

ag_ui_a2ui/toolkit/
prompt.rs

1//! Assembling the prompt for the surface-generating model.
2//!
3//! A2UI v0.9 is designed to be prompted rather than schema-constrained: the
4//! catalog goes into the prompt, the model writes JSON against it, and the
5//! result is validated afterwards. That makes prompt assembly load-bearing, so
6//! it lives here rather than being open-coded per agent.
7//!
8//! A prompt is assembled in this order:
9//!
10//! 1. the role,
11//! 2. the task,
12//! 3. `## Workflow Description:` — the rules that make output parseable,
13//! 4. `## UI Description:` — application-specific design guidance,
14//! 5. the catalog, as a compact summary or as the raw schema block,
15//! 6. `### Examples:` — few-shot examples, when supplied,
16//! 7. the conversation context,
17//! 8. the current surface, when updating one.
18//!
19//! Section 8 is what makes `intent = "update"` work: without the components
20//! already on screen, the model cannot edit them, only replace them.
21//!
22//! # Section headings are a cross-language contract
23//!
24//! The headings and the opening line of the workflow rules are the same strings
25//! every A2UI toolkit emits. Evaluation harnesses and prompt-diffing tools key
26//! on them, so a surface generated by this crate can be compared against one
27//! generated by the TypeScript or Python toolkit. The rules *under* that opening
28//! line are this crate's own, and are written to head off the failures its
29//! validator actually reports.
30
31use serde_json::Value;
32
33use crate::catalog::Catalog;
34use crate::constants::{A2UI_CLOSE_TAG, A2UI_OPEN_TAG, ROOT_ID};
35use crate::toolkit::history::PriorSurface;
36use crate::toolkit::ops::Intent;
37use crate::toolkit::schema::SchemaBundle;
38use crate::validate::ValidationError;
39
40/// Heading of the section carrying the generation rules.
41pub const WORKFLOW_SECTION: &str = "## Workflow Description:";
42/// Heading of the section carrying application design guidance.
43pub const UI_SECTION: &str = "## UI Description:";
44/// Heading of the section carrying few-shot examples.
45pub const EXAMPLES_SECTION: &str = "### Examples:";
46
47/// The rules that keep generated A2UI parseable and renderable.
48///
49/// Ordering is not cosmetic: a streaming renderer paints as components arrive,
50/// so `root` first and parents before children is what makes progressive
51/// rendering work.
52pub const GENERATION_GUIDELINES: &str = "\
53The generated response MUST follow these rules:
54- Reply with a flat list of A2UI messages as raw JSON. No prose inside the JSON.
55- Every component object is `{\"id\": <unique id>, \"component\": <type>, ...props}`.
56- Components are NEVER nested inline. A parent names its children by id.
57- Exactly one component must have `\"id\": \"root\"`; it is the top of the tree.
58- Order components top-down: `root` first, every parent before its children.
59- Reference only ids that exist in the same payload.
60- Never create a reference loop: a component must not be its own ancestor.
61- Static text goes in the component; only bind to the data model when the value \
62is dynamic.
63- Data bindings are `{\"path\": \"/pointer\"}` (JSON Pointer, absolute) and must \
64point at data you also send.
65- Relative paths (no leading `/`) are only valid inside a list template.
66- Use only component types and properties from the catalog below.";
67
68/// Everything the prompt builder needs.
69#[derive(Debug, Clone)]
70pub struct PromptSpec<'a> {
71    /// What the model is, e.g. "You generate UI surfaces for a travel agent."
72    pub role: &'a str,
73    /// What the user asked for on this turn.
74    pub request: &'a str,
75    /// Whether a new surface is being built or an existing one edited.
76    pub intent: Intent,
77    /// The surface id the model should target.
78    pub surface_id: &'a str,
79    /// The catalog the model must stay inside.
80    pub catalog: &'a Catalog,
81    /// The raw schema documents.
82    ///
83    /// When present the prompt carries the full JSON Schema block instead of the
84    /// compact catalog summary: exact, but far more tokens. Prune it first with
85    /// [`SchemaBundle::prune`] if the model only needs part of the catalog.
86    pub schemas: Option<&'a SchemaBundle>,
87    /// Few-shot examples, already rendered by
88    /// [`load_examples`](crate::toolkit::schema::load_examples).
89    pub examples: Option<&'a str>,
90    /// Prior conversation turns, oldest first, as `(role, text)`.
91    pub conversation: &'a [(String, String)],
92    /// The surface currently on screen, for [`Intent::Update`].
93    pub prior_surface: Option<&'a PriorSurface>,
94    /// Extra design guidance from the application.
95    pub ui_description: Option<&'a str>,
96    /// Replaces [`GENERATION_GUIDELINES`] wholesale.
97    ///
98    /// The default rules are tuned to what this crate's validator rejects; an
99    /// application with its own house style can substitute its own, keeping the
100    /// `The generated response MUST follow these rules:` opening line so the
101    /// section stays recognizable across toolkits.
102    pub workflow_rules: Option<&'a str>,
103    /// Whether to spell out the response format with the `<a2ui-json>` tags.
104    ///
105    /// Turn this off when the model answers through a structured-output tool,
106    /// where the tags would end up inside the JSON.
107    pub include_response_format: bool,
108}
109
110impl<'a> PromptSpec<'a> {
111    /// A spec with sensible defaults for a create-intent turn.
112    pub fn new(role: &'a str, request: &'a str, catalog: &'a Catalog) -> Self {
113        Self {
114            role,
115            request,
116            intent: Intent::Create,
117            surface_id: crate::constants::DEFAULT_SURFACE_ID,
118            catalog,
119            schemas: None,
120            examples: None,
121            conversation: &[],
122            prior_surface: None,
123            ui_description: None,
124            workflow_rules: None,
125            include_response_format: true,
126        }
127    }
128
129    /// Points the spec at an existing surface, switching to update intent.
130    #[must_use]
131    pub fn updating(mut self, prior: &'a PriorSurface) -> Self {
132        self.intent = Intent::Update;
133        self.surface_id = &prior.surface_id;
134        self.prior_surface = Some(prior);
135        self
136    }
137
138    /// Supplies prior conversation turns.
139    #[must_use]
140    pub fn with_conversation(mut self, conversation: &'a [(String, String)]) -> Self {
141        self.conversation = conversation;
142        self
143    }
144
145    /// Carries the raw schema block instead of the compact catalog summary.
146    #[must_use]
147    pub fn with_schemas(mut self, schemas: &'a SchemaBundle) -> Self {
148        self.schemas = Some(schemas);
149        self
150    }
151
152    /// Adds few-shot examples.
153    #[must_use]
154    pub fn with_examples(mut self, examples: &'a str) -> Self {
155        self.examples = Some(examples);
156        self
157    }
158}
159
160/// Builds the full system prompt for the generating model.
161pub fn build_subagent_prompt(spec: &PromptSpec<'_>) -> String {
162    let mut sections: Vec<String> = Vec::new();
163    sections.push(spec.role.trim().to_string());
164
165    sections.push(format!(
166        "## Task\n{}\n\nTarget surfaceId: {}\nIntent: {}",
167        spec.request.trim(),
168        spec.surface_id,
169        spec.intent
170    ));
171
172    let mut rules = spec
173        .workflow_rules
174        .unwrap_or(GENERATION_GUIDELINES)
175        .to_string();
176    if spec.intent == Intent::Update {
177        rules.push_str(&format!(
178            "\n- This surface already exists. Do NOT emit `createSurface` for '{}'; send only \
179             the components and data that change.",
180            spec.surface_id
181        ));
182    }
183    sections.push(format!("{WORKFLOW_SECTION}\n{rules}"));
184
185    if let Some(description) = spec.ui_description {
186        sections.push(format!("{UI_SECTION}\n{}", description.trim()));
187    }
188
189    // Either the exact schema documents or a compact summary of them, never
190    // both: the block is worth its tokens only when the model needs the detail.
191    match spec.schemas {
192        Some(schemas) => sections.push(schemas.render_llm_instructions()),
193        None => sections.push(spec.catalog.render_summary().trim_end().to_string()),
194    }
195
196    if let Some(examples) = spec.examples {
197        if !examples.trim().is_empty() {
198            sections.push(format!("{EXAMPLES_SECTION}\n{}", examples.trim_end()));
199        }
200    }
201
202    if !spec.conversation.is_empty() {
203        let mut block = String::from("### Conversation so far\n");
204        for (role, text) in spec.conversation {
205            block.push_str(&format!("{role}: {}\n", text.trim()));
206        }
207        sections.push(block.trim_end().to_string());
208    }
209
210    if let Some(prior) = spec.prior_surface {
211        sections.push(render_prior_surface(prior));
212    }
213
214    if spec.include_response_format {
215        sections.push(format!(
216            "## Response format\nWrap each A2UI JSON block in {A2UI_OPEN_TAG} and \
217             {A2UI_CLOSE_TAG}. Conversational text may go before or after a block, never inside \
218             one."
219        ));
220    }
221    sections.join("\n\n")
222}
223
224/// Renders the surface currently on screen, so the model can edit it.
225fn render_prior_surface(prior: &PriorSurface) -> String {
226    let components =
227        serde_json::to_string_pretty(&prior.components).unwrap_or_else(|_| "[]".to_string());
228    let data =
229        serde_json::to_string_pretty(&prior.data_model).unwrap_or_else(|_| "null".to_string());
230    let catalog = prior
231        .catalog_id
232        .as_deref()
233        .map(|id| format!("\ncatalogId: {id} (fixed; it cannot change)"))
234        .unwrap_or_default();
235
236    format!(
237        "## Surface currently on screen\nsurfaceId: {}{catalog}\n\nComponents:\n```json\n{components}\n```\n\n\
238         Data model:\n```json\n{data}\n```\n\nRe-send only the components you change, keeping \
239         their ids. Keep '{ROOT_ID}' as the root.",
240        prior.surface_id
241    )
242}
243
244/// Formats validation errors for a retry prompt.
245///
246/// One line per error, code first so the model can see the kind at a glance,
247/// then the locator, then the sentence explaining the fix.
248pub fn format_validation_errors(errors: &[ValidationError]) -> String {
249    if errors.is_empty() {
250        return String::new();
251    }
252    let mut out = String::from(
253        "The previous attempt was rejected. Fix every problem below and return the corrected \
254         A2UI, not a diff:\n",
255    );
256    for error in errors {
257        out.push_str(&format!(
258            "- [{}] at {}: {}\n",
259            error.code, error.path, error.message
260        ));
261    }
262    out.trim_end().to_string()
263}
264
265/// Appends the errors from a failed attempt to the prompt for the next one.
266pub fn augment_prompt_with_errors(prompt: &str, errors: &[ValidationError]) -> String {
267    if errors.is_empty() {
268        return prompt.to_string();
269    }
270    format!(
271        "{prompt}\n\n## Correction required\n{}",
272        format_validation_errors(errors)
273    )
274}
275
276/// A short human-readable summary of a surface, for logs and activity reports.
277pub fn describe_surface(prior: &PriorSurface) -> String {
278    let mut kinds: Vec<&str> = prior
279        .components
280        .iter()
281        .map(|c| c.component.as_str())
282        .collect();
283    kinds.sort_unstable();
284    kinds.dedup();
285    format!(
286        "surface '{}' with {} component(s) [{}]",
287        prior.surface_id,
288        prior.components.len(),
289        kinds.join(", ")
290    )
291}
292
293/// Renders a data model compactly for inclusion in a prompt.
294pub fn render_data_model(data_model: &Value) -> String {
295    serde_json::to_string_pretty(data_model).unwrap_or_else(|_| "null".to_string())
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use crate::message::Component;
302    use serde_json::json;
303
304    fn prior() -> PriorSurface {
305        PriorSurface {
306            surface_id: "cart".to_string(),
307            catalog_id: Some("cat".to_string()),
308            components: vec![
309                Component::new("root", "Column").with("children", json!(["title"])),
310                Component::new("title", "Text").with("text", json!("Cart")),
311            ],
312            data_model: json!({"total": 12}),
313            deleted: false,
314        }
315    }
316
317    #[test]
318    fn a_create_prompt_has_every_section() {
319        let catalog = Catalog::basic();
320        let conversation = vec![("user".to_string(), "show my cart".to_string())];
321        let spec = PromptSpec::new("You build UI.", "Render the cart.", &catalog)
322            .with_conversation(&conversation);
323        let prompt = build_subagent_prompt(&spec);
324
325        assert!(prompt.starts_with("You build UI."));
326        assert!(prompt.contains("## Task"));
327        assert!(prompt.contains("Intent: create"));
328        assert!(prompt.contains("## Workflow Description:"));
329        assert!(prompt.contains("The generated response MUST follow these rules:"));
330        assert!(prompt.contains("### Component catalog"));
331        assert!(prompt.contains("- Text required: text (value)"));
332        assert!(prompt.contains("### Conversation so far"));
333        assert!(prompt.contains("user: show my cart"));
334        assert!(prompt.contains(A2UI_OPEN_TAG));
335        assert!(!prompt.contains("## Surface currently on screen"));
336    }
337
338    #[test]
339    fn an_update_prompt_forbids_create_surface_and_shows_the_current_tree() {
340        let catalog = Catalog::basic();
341        let prior = prior();
342        let spec =
343            PromptSpec::new("You build UI.", "Add a checkout button.", &catalog).updating(&prior);
344        let prompt = build_subagent_prompt(&spec);
345
346        assert!(prompt.contains("Intent: update"));
347        assert!(prompt.contains("Do NOT emit `createSurface` for 'cart'"));
348        assert!(prompt.contains("## Surface currently on screen"));
349        assert!(prompt.contains("catalogId: cat (fixed; it cannot change)"));
350        assert!(prompt.contains("\"id\": \"title\""));
351        assert!(prompt.contains("\"total\": 12"));
352        assert!(prompt.contains("Target surfaceId: cart"));
353    }
354
355    #[test]
356    fn the_response_format_section_can_be_suppressed() {
357        let catalog = Catalog::basic();
358        let mut spec = PromptSpec::new("role", "request", &catalog);
359        spec.include_response_format = false;
360        let prompt = build_subagent_prompt(&spec);
361        assert!(!prompt.contains(A2UI_OPEN_TAG));
362    }
363
364    #[test]
365    fn design_guidance_is_included_when_supplied() {
366        let catalog = Catalog::basic();
367        let mut spec = PromptSpec::new("role", "request", &catalog);
368        spec.ui_description = Some("Use the brand blue for primary buttons.");
369        let prompt = build_subagent_prompt(&spec);
370        assert!(prompt.contains("## UI Description:"));
371        assert!(prompt.contains("Use the brand blue"));
372    }
373
374    #[test]
375    fn validation_errors_render_one_actionable_line_each() {
376        let errors = vec![
377            ValidationError::new(
378                crate::validate::ErrorCode::NoRoot,
379                "components",
380                "No component has id 'root'.",
381            ),
382            ValidationError::new(
383                crate::validate::ErrorCode::ChildCycle,
384                "components[1].child",
385                "Child references form a loop: a -> b -> a.",
386            ),
387        ];
388        let rendered = format_validation_errors(&errors);
389        assert!(rendered.contains("- [no_root] at components: No component has id 'root'."));
390        assert!(rendered.contains("- [child_cycle] at components[1].child:"));
391        assert_eq!(format_validation_errors(&[]), "");
392
393        let augmented = augment_prompt_with_errors("base prompt", &errors);
394        assert!(augmented.starts_with("base prompt"));
395        assert!(augmented.contains("## Correction required"));
396        assert_eq!(
397            augment_prompt_with_errors("base prompt", &[]),
398            "base prompt"
399        );
400    }
401
402    #[test]
403    fn surfaces_summarize_for_logs() {
404        assert_eq!(
405            describe_surface(&prior()),
406            "surface 'cart' with 2 component(s) [Column, Text]"
407        );
408    }
409
410    #[test]
411    fn guidelines_state_the_rules_the_validator_enforces() {
412        for rule in [
413            "\"id\": \"root\"",
414            "NEVER nested inline",
415            "reference loop",
416            "top-down",
417        ] {
418            assert!(
419                GENERATION_GUIDELINES
420                    .to_lowercase()
421                    .contains(&rule.to_lowercase()),
422                "guidelines omit {rule}"
423            );
424        }
425    }
426}