Skip to main content

talos_agent/prompt/
builder.rs

1use std::collections::BTreeMap;
2use std::collections::HashMap;
3
4use talos_core::message::{SystemCacheMarker, SystemCacheType};
5use talos_plugin::{HookContext, HookEvent, HookOutcome, HookRegistry};
6use talos_skill::SkillIndex;
7
8use super::assets::{DEFAULT_IDENTITY, TOOL_CALLING_STRICT};
9use super::sections::{PromptSection, PromptSectionKind};
10use super::types::{ActivatedSkillContext, CacheMarker, CacheType, ContextFile, ToolDescription};
11
12/// Builder for assembling a system prompt from multiple components.
13///
14/// The builder uses a fluent API to configure each component of the system
15/// prompt. Components are assembled in a fixed order optimized for LLM
16/// provider caching: stable sections first, then semi-stable, then dynamic.
17///
18/// # Component Order
19///
20/// 1. Identity (or custom prompt if provided)
21/// 2. Tool descriptions
22/// 3. Skill index (Level 0)
23/// 4. Context files (AGENTS.md)
24/// 5. User preferences
25/// 6. Append prompt (if provided)
26///
27/// # Example
28///
29/// ```
30/// use talos_agent::prompt::SystemPromptBuilder;
31///
32/// let prompt = SystemPromptBuilder::new()
33///     .with_user_preferences("Always use British English.".into())
34///     .build();
35/// ```
36#[derive(Debug, Clone)]
37pub struct SystemPromptBuilder {
38    /// Agent identity and role instructions.
39    identity: String,
40    /// Tool names and descriptions.
41    tools: Vec<ToolDescription>,
42    /// Level 0 skill index (name + description).
43    skill_index: Vec<SkillIndex>,
44    /// Explicitly activated Level 1/2 Skill content.
45    activated_skill: Option<ActivatedSkillContext>,
46    /// AGENTS.md and other context file contents.
47    context_files: Vec<ContextFile>,
48    /// User-specific instructions.
49    user_preferences: String,
50    /// Overrides the default identity when provided.
51    custom_prompt: Option<String>,
52    /// Appended to the end of the prompt when provided.
53    append_prompt: Option<String>,
54    /// Bounded memory injection section (advisory, never authoritative).
55    memory_section: Option<String>,
56    /// Bounded session todo section (advisory orchestration context).
57    todo_section: Option<String>,
58    tool_call_format: &'static str,
59    /// Runtime template values used for `{{slot}}` substitution.
60    template_vars: HashMap<String, String>,
61}
62
63impl SystemPromptBuilder {
64    /// Creates a new builder with the default identity and no other components.
65    ///
66    /// All optional components start empty. Use the builder methods to
67    /// configure tools, skills, context files, and other components.
68    #[must_use]
69    pub fn new() -> Self {
70        let mut template_vars = HashMap::new();
71        template_vars.insert(
72            "workspace_info".to_string(),
73            "Workspace information unavailable.".to_string(),
74        );
75        template_vars.insert(
76            "model_info".to_string(),
77            "Provider model metadata unavailable.".to_string(),
78        );
79
80        Self {
81            identity: DEFAULT_IDENTITY.to_string(),
82            tools: Vec::new(),
83            skill_index: Vec::new(),
84            activated_skill: None,
85            context_files: Vec::new(),
86            user_preferences: String::new(),
87            custom_prompt: None,
88            append_prompt: None,
89            memory_section: None,
90            todo_section: None,
91            tool_call_format: "",
92            template_vars,
93        }
94    }
95
96    pub fn with_strict_tool_format(mut self) -> Self {
97        self.tool_call_format = TOOL_CALLING_STRICT;
98        self
99    }
100
101    pub fn with_tool_format(mut self, format: &'static str) -> Self {
102        self.tool_call_format = format;
103        self
104    }
105
106    /// Sets a template slot value for `{{slot}}` substitution.
107    #[must_use]
108    pub fn with_template_var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
109        self.template_vars.insert(key.into(), value.into());
110        self
111    }
112
113    /// Sets workspace information for the default identity template.
114    #[must_use]
115    pub fn with_workspace_info(self, value: impl Into<String>) -> Self {
116        self.with_template_var("workspace_info", value)
117    }
118
119    /// Sets model information for the default identity template.
120    #[must_use]
121    pub fn with_model_info(self, value: impl Into<String>) -> Self {
122        self.with_template_var("model_info", value)
123    }
124
125    /// Sets the tool descriptions for inclusion in the system prompt.
126    ///
127    /// Tools are sorted alphabetically by name to ensure stable ordering
128    /// across turns, maximizing cache hit rates.
129    #[must_use]
130    pub fn with_tools(mut self, tools: Vec<ToolDescription>) -> Self {
131        self.tools = tools;
132        self
133    }
134
135    /// Sets the skill index for inclusion in the system prompt.
136    ///
137    /// Only Level 0 metadata (name, description, triggers) is included.
138    /// Full skill bodies are not loaded at this stage.
139    #[must_use]
140    pub fn with_skill_index(mut self, skills: Vec<SkillIndex>) -> Self {
141        self.skill_index = skills;
142        self
143    }
144
145    /// Sets the explicitly activated Skill context.
146    ///
147    /// This content is cacheable after activation. The owning [`Agent`](crate::Agent)
148    /// invalidates the stable-prefix cache whenever activation changes.
149    #[must_use]
150    pub fn with_activated_skill(mut self, skill: Option<ActivatedSkillContext>) -> Self {
151        self.activated_skill = skill;
152        self
153    }
154
155    /// Sets the context files for inclusion in the system prompt.
156    ///
157    /// Typically loaded from `AGENTS.md` files via [`ContextLoader`].
158    ///
159    /// [`ContextLoader`]: crate::context::ContextLoader
160    #[must_use]
161    pub fn with_context_files(mut self, files: Vec<ContextFile>) -> Self {
162        self.context_files = files;
163        self
164    }
165
166    /// Sets user-specific instructions for inclusion in the system prompt.
167    #[must_use]
168    pub fn with_user_preferences(mut self, prefs: String) -> Self {
169        self.user_preferences = prefs;
170        self
171    }
172
173    /// Sets a bounded memory section for inclusion in the system prompt.
174    ///
175    /// Memory is advisory only — never authoritative over session context.
176    /// When `None`, no memory section is injected.
177    #[must_use]
178    pub fn with_memory_section(mut self, section: Option<String>) -> Self {
179        self.memory_section = section;
180        self
181    }
182
183    /// Sets a bounded session todo section for inclusion in the dynamic prompt suffix.
184    ///
185    /// Todo context is advisory orchestration state. It is intentionally not part of the
186    /// stable cached prefix because it can change between turns.
187    #[must_use]
188    pub fn with_todo_section(mut self, section: Option<String>) -> Self {
189        self.todo_section = section;
190        self
191    }
192
193    /// Sets a custom prompt that replaces the default identity.
194    ///
195    /// When provided, the custom prompt is used instead of the default
196    /// identity. The rest of the prompt (tools, skills, etc.) is still
197    /// assembled normally.
198    #[must_use]
199    pub fn with_custom_prompt(mut self, prompt: String) -> Self {
200        self.custom_prompt = Some(prompt);
201        self
202    }
203
204    /// Sets an append prompt that is added at the end of the system prompt.
205    ///
206    /// The append prompt is always placed last, after all other components.
207    #[must_use]
208    pub fn with_append_prompt(mut self, prompt: String) -> Self {
209        self.append_prompt = Some(prompt);
210        self
211    }
212
213    /// Clears the append prompt, removing any previously set value.
214    pub fn clear_append_prompt(&mut self) {
215        self.append_prompt = None;
216    }
217
218    /// Sets the append prompt to an optional value.
219    ///
220    /// Use `None` to clear the append prompt, or `Some(prompt)` to set it.
221    pub fn set_append_prompt_opt(&mut self, prompt: Option<String>) {
222        self.append_prompt = prompt;
223    }
224
225    fn render_template(&self, template: &str, extra_vars: &[(&str, String)]) -> String {
226        let mut rendered = template.to_string();
227        let mut vars = self.template_vars.clone();
228        for (key, value) in extra_vars {
229            vars.insert((*key).to_string(), value.clone());
230        }
231
232        for (key, value) in vars {
233            rendered = rendered.replace(&format!("{{{{{key}}}}}"), &value);
234        }
235        rendered
236    }
237
238    fn current_datetime() -> String {
239        let seconds = std::time::SystemTime::now()
240            .duration_since(std::time::UNIX_EPOCH)
241            .map(|duration| duration.as_secs())
242            .unwrap_or(0);
243        format!("unix_seconds={seconds}")
244    }
245
246    fn tool_protocol_hint(&self) -> String {
247        if self.tool_call_format.is_empty() {
248            "Native tool calling is enabled. Use provider-native tool calls; do not emit textual tool-call JSON unless the provider requires a fallback.".to_string()
249        } else {
250            self.tool_call_format.trim().to_string()
251        }
252    }
253
254    fn prompt_sections(&self) -> Vec<PromptSection> {
255        let mut sections: Vec<PromptSection> = Vec::new();
256
257        let stable_vars = [("tool_protocol_hint", self.tool_protocol_hint())];
258
259        let identity = if let Some(ref custom) = self.custom_prompt {
260            self.render_template(custom, &stable_vars)
261        } else {
262            self.render_template(&self.identity, &stable_vars)
263        };
264        sections.push(PromptSection {
265            text: format!("# Identity\n{identity}\n"),
266            kind: PromptSectionKind::Cacheable,
267        });
268
269        if self.tools.is_empty() {
270            sections.push(PromptSection {
271                text: String::from("# Tools\nNo tools available.\n"),
272                kind: PromptSectionKind::Cacheable,
273            });
274        } else {
275            let mut families: BTreeMap<_, Vec<&ToolDescription>> = BTreeMap::new();
276            for tool in &self.tools {
277                families.entry(tool.family).or_default().push(tool);
278            }
279
280            sections.push(PromptSection {
281                text: String::from("# Tools\nTool definitions are grouped by stable family.\n"),
282                kind: PromptSectionKind::Cacheable,
283            });
284
285            for (family, mut sorted_tools) in families {
286                sorted_tools.sort_by(|a, b| a.name.cmp(&b.name));
287                let mut tools_section = format!("# Tool Family: {family:?}\n");
288                for tool in sorted_tools {
289                    tools_section.push_str(&format!("## {}\n{}\n", tool.name, tool.description));
290                    if let Some(props) = tool.parameters.get("properties")
291                        && let Some(required) = tool.parameters.get("required")
292                    {
293                        let req_list: Vec<&str> = required
294                            .as_array()
295                            .map(|a| a.iter().filter_map(|v| v.as_str()).collect())
296                            .unwrap_or_default();
297                        let mut param_parts = Vec::new();
298                        for (key, val) in props.as_object().unwrap_or(&serde_json::Map::new()) {
299                            let desc = val
300                                .get("description")
301                                .and_then(|d| d.as_str())
302                                .unwrap_or("");
303                            let ptype = val.get("type").and_then(|t| t.as_str()).unwrap_or("any");
304                            let req = if req_list.contains(&key.as_str()) {
305                                "required"
306                            } else {
307                                "optional"
308                            };
309                            param_parts
310                                .push(format!("  - {} ({}): {} [{}]", key, ptype, desc, req));
311                        }
312                        if !param_parts.is_empty() {
313                            tools_section.push_str("Parameters:\n");
314                            tools_section.push_str(&param_parts.join("\n"));
315                            tools_section.push_str("\n\n");
316                        }
317                    }
318                    tools_section.push('\n');
319                }
320                sections.push(PromptSection {
321                    text: tools_section,
322                    kind: PromptSectionKind::Cacheable,
323                });
324            }
325        }
326
327        if self.skill_index.is_empty() {
328            sections.push(PromptSection {
329                text: String::from("# Skills\nNo skills available.\n"),
330                kind: PromptSectionKind::Cacheable,
331            });
332        } else {
333            let mut skills_section = String::from("# Skills\n");
334            for skill in &self.skill_index {
335                skills_section.push_str(&format!("- **{}**: {}\n", skill.name, skill.description));
336            }
337            skills_section.push('\n');
338            sections.push(PromptSection {
339                text: skills_section,
340                kind: PromptSectionKind::Cacheable,
341            });
342        }
343
344        if let Some(ref skill) = self.activated_skill {
345            sections.push(PromptSection {
346                text: format!(
347                    "# Activated Skill: {}\n{}\n",
348                    skill.name.trim(),
349                    skill.content.trim()
350                ),
351                kind: PromptSectionKind::Cacheable,
352            });
353        }
354
355        if self.context_files.is_empty() {
356            sections.push(PromptSection {
357                text: String::from("# Context\nNo context files loaded.\n"),
358                kind: PromptSectionKind::Dynamic,
359            });
360        } else {
361            let mut context_section = String::from("# Context\n");
362            for file in &self.context_files {
363                context_section.push_str(&format!("--- {} ---\n{}\n\n", file.path, file.content));
364            }
365            sections.push(PromptSection {
366                text: context_section,
367                kind: PromptSectionKind::Dynamic,
368            });
369        }
370
371        if let Some(ref memory) = self.memory_section {
372            sections.push(PromptSection {
373                text: format!("# Memory\n{memory}\n"),
374                kind: PromptSectionKind::Dynamic,
375            });
376        }
377
378        if let Some(ref todos) = self.todo_section {
379            sections.push(PromptSection {
380                text: format!("# Session Todos\n{todos}\n"),
381                kind: PromptSectionKind::Dynamic,
382            });
383        }
384
385        if !self.user_preferences.is_empty() {
386            sections.push(PromptSection {
387                text: format!("# User Preferences\n{}\n", self.user_preferences),
388                kind: PromptSectionKind::Dynamic,
389            });
390        }
391
392        let runtime_section = self.render_template(
393            "# Runtime Context\nCurrent datetime: {{datetime}}\n",
394            &[("datetime", Self::current_datetime())],
395        );
396        sections.push(PromptSection {
397            text: runtime_section,
398            kind: PromptSectionKind::Dynamic,
399        });
400
401        if let Some(ref append) = self.append_prompt {
402            sections.push(PromptSection {
403                text: format!("# Additional Instructions\n{append}\n"),
404                kind: PromptSectionKind::Dynamic,
405            });
406        }
407
408        sections
409    }
410
411    /// Assembles and returns the final system prompt as a string.
412    ///
413    /// Components are assembled in the optimal order for caching:
414    /// 1. Identity (or custom prompt if provided)
415    /// 2. Tools (sorted by name)
416    /// 3. Skill index
417    /// 4. Context files
418    /// 5. User preferences
419    /// 6. Append prompt (if provided)
420    ///
421    /// Empty components are omitted from the output.
422    #[must_use]
423    pub fn build(&self) -> String {
424        self.prompt_sections()
425            .into_iter()
426            .map(|section| section.text)
427            .collect::<Vec<_>>()
428            .join("\n")
429    }
430
431    /// Assembles the system prompt and emits the `OnSystemPromptBuilt` hook.
432    ///
433    /// Handlers may replace the final prompt by returning
434    /// [`talos_plugin::HookResult::Modify`]. `Skip` leaves the prompt unchanged.
435    #[allow(dead_code)]
436    pub(crate) async fn build_with_hooks(
437        &self,
438        hook_registry: &HookRegistry,
439        ctx: &HookContext,
440    ) -> Result<(String, Vec<SystemCacheMarker>), String> {
441        let (prompt, markers) = self.build_with_cache_markers();
442        let original_prompt = prompt.clone();
443        let outcome = hook_registry
444            .dispatch(ctx, HookEvent::OnSystemPromptBuilt { prompt: &prompt })
445            .await;
446
447        match outcome {
448            HookOutcome::Continue(HookEvent::OnSystemPromptBuilt { prompt })
449            | HookOutcome::Skip(HookEvent::OnSystemPromptBuilt { prompt }) => {
450                let prompt = prompt.to_string();
451                let markers = if prompt == original_prompt {
452                    markers.into_iter().map(Into::into).collect()
453                } else {
454                    Vec::new()
455                };
456                Ok((prompt, markers))
457            }
458            HookOutcome::Deny { reason, .. } => Err(reason),
459            HookOutcome::Continue(_) | HookOutcome::Skip(_) => {
460                Ok((prompt, markers.into_iter().map(Into::into).collect()))
461            }
462        }
463    }
464
465    /// Runs a pre-assembled prompt through the `OnSystemPromptBuilt` hook.
466    ///
467    /// The `stable_prefix_len` indicates the byte length of the stable prefix
468    /// (Identity + Tools + Skills) within the combined prompt. A cache marker
469    /// is emitted for this range if the hook does not modify the prompt.
470    pub(crate) async fn build_with_hooks_from_prompt(
471        &self,
472        hook_registry: &HookRegistry,
473        ctx: &HookContext,
474        prompt: &str,
475        stable_prefix_len: usize,
476    ) -> Result<(String, Vec<SystemCacheMarker>), String> {
477        let original_prompt = prompt.to_string();
478        let outcome = hook_registry
479            .dispatch(ctx, HookEvent::OnSystemPromptBuilt { prompt })
480            .await;
481
482        match outcome {
483            HookOutcome::Continue(HookEvent::OnSystemPromptBuilt { prompt })
484            | HookOutcome::Skip(HookEvent::OnSystemPromptBuilt { prompt }) => {
485                let prompt = prompt.to_string();
486                let markers = if prompt == original_prompt && stable_prefix_len > 0 {
487                    vec![SystemCacheMarker {
488                        offset: 0,
489                        length: stable_prefix_len,
490                        cache_type: SystemCacheType::Ephemeral,
491                    }]
492                } else {
493                    Vec::new()
494                };
495                Ok((prompt, markers))
496            }
497            HookOutcome::Deny { reason, .. } => Err(reason),
498            HookOutcome::Continue(_) | HookOutcome::Skip(_) => {
499                let markers = if stable_prefix_len > 0 {
500                    vec![SystemCacheMarker {
501                        offset: 0,
502                        length: stable_prefix_len,
503                        cache_type: SystemCacheType::Ephemeral,
504                    }]
505                } else {
506                    Vec::new()
507                };
508                Ok((prompt.to_string(), markers))
509            }
510        }
511    }
512
513    /// Builds only the stable prefix (Identity + Tools + Skills).
514    ///
515    /// These sections are cacheable and do not change between turns unless
516    /// tools, skills, or the identity/custom prompt are modified. The result
517    /// can be cached by the caller and reused across turns.
518    ///
519    /// Returns `None` if there are no stable sections (should not happen in
520    /// practice since Identity is always present).
521    #[must_use]
522    pub fn build_stable_prefix(&self) -> String {
523        let sections = self.prompt_sections();
524        let mut prefix = String::new();
525        let mut first = true;
526        for section in &sections {
527            if section.kind != PromptSectionKind::Cacheable {
528                break;
529            }
530            if !first {
531                prefix.push('\n');
532            }
533            prefix.push_str(&section.text);
534            first = false;
535        }
536        prefix
537    }
538
539    /// Builds only the dynamic suffix (Context + User Preferences + Runtime + Append).
540    ///
541    /// These sections change every turn (e.g., datetime) or are semi-stable
542    /// (context files, user preferences). Combined with a cached stable prefix,
543    /// they form the complete system prompt.
544    #[must_use]
545    pub fn build_dynamic_suffix(&self) -> String {
546        let sections = self.prompt_sections();
547        let mut suffix = String::new();
548        let mut first = true;
549        for section in &sections {
550            if section.kind == PromptSectionKind::Cacheable {
551                continue;
552            }
553            if !first {
554                suffix.push('\n');
555            }
556            suffix.push_str(&section.text);
557            first = false;
558        }
559        suffix
560    }
561
562    /// Assembles the system prompt with cache control markers.
563    ///
564    /// Returns the prompt string and a list of [`CacheMarker`]s indicating
565    /// which byte ranges are stable and suitable for provider caching.
566    ///
567    /// Cacheable sections:
568    /// - Identity (or custom prompt)
569    /// - Tools
570    /// - Skill index
571    ///
572    /// Semi-stable sections (context files, user preferences) and the
573    /// append prompt are not marked for caching.
574    #[must_use]
575    pub fn build_with_cache_markers(&self) -> (String, Vec<CacheMarker>) {
576        let mut markers: Vec<CacheMarker> = Vec::new();
577        let sections = self.prompt_sections();
578        let mut prompt = String::new();
579
580        for (index, section) in sections.iter().enumerate() {
581            if index > 0 {
582                prompt.push('\n');
583            }
584            let offset = prompt.len();
585            prompt.push_str(&section.text);
586            if section.kind == PromptSectionKind::Cacheable {
587                markers.push(CacheMarker {
588                    offset,
589                    length: section.text.len(),
590                    cache_type: CacheType::Ephemeral,
591                });
592            }
593        }
594
595        (prompt, markers)
596    }
597
598    /// Estimates the total token count of the assembled prompt.
599    ///
600    /// Uses a heuristic of 1 token per 4 characters, which is a reasonable
601    /// approximation for English text. This is not exact and should not be
602    /// used for billing purposes.
603    #[must_use]
604    pub fn total_tokens(&self) -> usize {
605        let prompt = self.build();
606        // Heuristic: ~1 token per 4 characters for English text
607        prompt.chars().count().div_ceil(4)
608    }
609
610    /// Logs the prompt size for debugging purposes.
611    ///
612    /// Prints the character count and estimated token count to stderr.
613    /// This is useful for monitoring prompt size during development.
614    pub fn log_size(&self) {
615        let prompt = self.build();
616        let char_count = prompt.chars().count();
617        let token_estimate = self.total_tokens();
618        eprintln!("System prompt: {char_count} characters, ~{token_estimate} tokens");
619    }
620}
621
622impl Default for SystemPromptBuilder {
623    fn default() -> Self {
624        Self::new()
625    }
626}