telos_agent/agent/prompt/
assembly.rs

1use crate::agent::prompt::PromptSection;
2use crate::agent::prompt::section::{PromptBlock, PromptStability};
3use std::collections::HashMap;
4use tokio::sync::Mutex;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct PromptSectionStat {
8    pub name: String,
9    pub chars: usize,
10    pub stability: PromptStability,
11}
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct PromptStats {
15    pub total_chars: usize,
16    pub sections: Vec<PromptSectionStat>,
17}
18
19/// Assembles a system prompt from ordered sections with caching.
20/// Static sections are rendered once and cached; dynamic re-render each time.
21pub struct PromptAssembly {
22    sections: Vec<Box<dyn PromptSection>>,
23    static_cache: Mutex<HashMap<String, String>>,
24}
25
26impl Default for PromptAssembly {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl PromptAssembly {
33    pub fn new() -> Self {
34        Self { sections: Vec::new(), static_cache: Mutex::new(HashMap::new()) }
35    }
36
37    /// Add a section whose stability is determined by its own
38    /// [`PromptSection::stability`] implementation.
39    pub fn add(&mut self, section: impl PromptSection + 'static) {
40        self.sections.push(Box::new(section));
41    }
42
43    pub async fn build(&self) -> String {
44        let mut parts = Vec::new();
45        for section in &self.sections {
46            let text = self.render_section(section).await;
47            if !text.is_empty() {
48                parts.push(text);
49            }
50        }
51        parts.join("\n\n")
52    }
53
54    /// Render the assembly into structured blocks.
55    pub async fn build_blocks(&self) -> Vec<PromptBlock> {
56        let mut blocks = Vec::new();
57        for section in &self.sections {
58            let text = self.render_section(section).await;
59            if !text.is_empty() {
60                blocks.push(PromptBlock {
61                    name: section.name().to_string(),
62                    text,
63                    stability: section.stability(),
64                });
65            }
66        }
67        blocks
68    }
69
70    pub async fn build_stats(&self) -> PromptStats {
71        let blocks = self.build_blocks().await;
72        let sections = blocks
73            .iter()
74            .map(|block| PromptSectionStat {
75                name: block.name.clone(),
76                chars: block.text.chars().count(),
77                stability: block.stability,
78            })
79            .collect::<Vec<_>>();
80        let total_chars = sections.iter().map(|section| section.chars).sum();
81        PromptStats { total_chars, sections }
82    }
83
84    async fn render_section(&self, section: &dyn PromptSection) -> String {
85        match section.stability() {
86            PromptStability::Static => {
87                let mut cache = self.static_cache.lock().await;
88                if let Some(cached) = cache.get(section.name()) {
89                    cached.clone()
90                } else {
91                    let rendered = section.render(&()).await;
92                    cache.insert(section.name().to_string(), rendered.clone());
93                    rendered
94                }
95            }
96            PromptStability::Dynamic => section.render(&()).await,
97        }
98    }
99}
telos_agent/agent/prompt/assembly.rs

telos_agent/agent/prompt/
assembly.rs
