ai_agent/constants/
system_prompt_sections.rs

1//! System prompt sections module
2//!
3//! Provides functions for creating and managing system prompt sections
4//! that can be cached or computed on each turn.
5
6use std::collections::HashMap;
7
8/// A function that computes a prompt section value
9pub type ComputeFn = Box<dyn Fn() -> Option<String> + Send + Sync>;
10
11/// A system prompt section with name, compute function, and cache behavior
12pub struct SystemPromptSection {
13    pub name: String,
14    pub compute: ComputeFn,
15    pub cache_break: bool,
16}
17
18/// Create a memoized system prompt section.
19/// Computed once, cached until /clear or /compact.
20pub fn system_prompt_section(name: &str, compute: ComputeFn) -> SystemPromptSection {
21    SystemPromptSection {
22        name: name.to_string(),
23        compute,
24        cache_break: false,
25    }
26}
27
28/// Create a volatile system prompt section that recomputes every turn.
29/// This WILL break the prompt cache when the value changes.
30/// Requires a reason explaining why cache-breaking is necessary.
31pub fn dangerous_uncached_system_prompt_section(
32    name: &str,
33    compute: ComputeFn,
34    _reason: &str,
35) -> SystemPromptSection {
36    SystemPromptSection {
37        name: name.to_string(),
38        compute,
39        cache_break: true,
40    }
41}
42
43/// Resolve all system prompt sections, returning prompt strings.
44/// Uses a cache to avoid recomputing sections unnecessarily.
45pub fn resolve_system_prompt_sections(
46    sections: &[SystemPromptSection],
47    cache: &mut HashMap<String, Option<String>>,
48) -> Vec<Option<String>> {
49    sections
50        .iter()
51        .map(|s| {
52            if !s.cache_break {
53                if let Some(cached) = cache.get(&s.name) {
54                    return cached.clone();
55                }
56            }
57            let value = (s.compute)();
58            cache.insert(s.name.clone(), value.clone());
59            value
60        })
61        .collect()
62}
63
64/// Clear all system prompt section state.
65/// Called on /clear and /compact.
66pub fn clear_system_prompt_sections(_cache: &mut HashMap<String, Option<String>>) {
67    // Clear the cache
68    _cache.clear();
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn test_system_prompt_section() {
77        let compute = Box::new(|| Some("test prompt".to_string()));
78        let section = system_prompt_section("test", compute);
79
80        assert_eq!(section.name, "test");
81        assert!(!section.cache_break);
82    }
83
84    #[test]
85    fn test_uncached_section() {
86        let compute = Box::new(|| Some("test prompt".to_string()));
87        let section =
88            dangerous_uncached_system_prompt_section("test", compute, "needs fresh value");
89
90        assert_eq!(section.name, "test");
91        assert!(section.cache_break);
92    }
93
94    #[test]
95    fn test_resolve_with_cache() {
96        let compute = Box::new(|| Some("computed value".to_string()));
97        let section = system_prompt_section("test", compute);
98
99        let mut cache = HashMap::new();
100        let results = resolve_system_prompt_sections(&[section], &mut cache);
101
102        assert_eq!(results.len(), 1);
103        assert_eq!(results[0], Some("computed value".to_string()));
104    }
105
106    #[test]
107    fn test_cache_hit() {
108        let compute = Box::new(|| Some("new value".to_string()));
109        let section = system_prompt_section("test", compute);
110
111        let mut cache = HashMap::new();
112        cache.insert("test".to_string(), Some("cached value".to_string()));
113
114        let results = resolve_system_prompt_sections(&[section], &mut cache);
115
116        assert_eq!(results[0], Some("cached value".to_string()));
117    }
118}
ai_agent/constants/system_prompt_sections.rs

ai_agent/constants/
system_prompt_sections.rs
