matrixcode_core/prompt/
orchestrator.rs

1//! Prompt Orchestrator
2//!
3//! Core prompt assembly system that:
4//! - Composes sections in order
5//! - Inserts cache boundary markers
6//! - Manages caching for static sections
7//! - Injects runtime context
8
9use std::sync::Arc;
10use crate::prompt::{PromptSection, SectionCache, CacheKey};
11use crate::prompt::{ContextInjector, UserContext, SystemContext};
12
13/// Cache boundary marker for API caching
14/// 
15/// This marker indicates where cached content ends and dynamic content begins.
16/// APIs like Claude can use this for prompt prefix caching.
17pub const CACHE_BOUNDARY: &str = "\n<!-- CACHE_BOUNDARY -->\n";
18
19/// Prompt profile type
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
21pub enum PromptProfile {
22    /// Default profile (full capabilities)
23    Default,
24    /// Safe profile (restricted operations)
25    Safe,
26    /// Fast profile (minimal prompt)
27    Fast,
28    /// Review profile (code review focus)
29    Review,
30}
31
32impl PromptProfile {
33    pub fn as_str(&self) -> &'static str {
34        match self {
35            Self::Default => "default",
36            Self::Safe => "safe",
37            Self::Fast => "fast",
38            Self::Review => "review",
39        }
40    }
41
42    pub fn from_str(s: &str) -> Self {
43        match s {
44            "safe" => Self::Safe,
45            "fast" => Self::Fast,
46            "review" => Self::Review,
47            _ => Self::Default,
48        }
49    }
50}
51
52impl std::fmt::Display for PromptProfile {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        write!(f, "{}", self.as_str())
55    }
56}
57
58impl Default for PromptProfile {
59    fn default() -> Self {
60        Self::Default
61    }
62}
63
64/// Prompt orchestrator - manages prompt assembly
65pub struct PromptOrchestrator {
66    /// Cache for static sections
67    cache: Arc<SectionCache>,
68    /// Context injector
69    context_injector: ContextInjector,
70    /// Current profile
71    profile: PromptProfile,
72    /// Sections to include
73    sections: Vec<PromptSection>,
74    /// Whether to include cache boundary
75    include_boundary: bool,
76    /// Whether to inject context
77    inject_context: bool,
78}
79
80impl PromptOrchestrator {
81    /// Create a new orchestrator
82    pub fn new<P: Into<std::path::PathBuf>>(working_dir: P) -> Self {
83        Self {
84            cache: Arc::new(SectionCache::new()),
85            context_injector: ContextInjector::new(working_dir.into()),
86            profile: PromptProfile::Default,
87            sections: Vec::new(),
88            include_boundary: true,
89            inject_context: true,
90        }
91    }
92
93    /// Create with shared cache
94    pub fn with_cache(mut self, cache: Arc<SectionCache>) -> Self {
95        self.cache = cache;
96        self
97    }
98
99    /// Set profile
100    pub fn with_profile(mut self, profile: PromptProfile) -> Self {
101        self.profile = profile;
102        self
103    }
104
105    /// Set whether to include cache boundary
106    pub fn with_boundary(mut self, include: bool) -> Self {
107        self.include_boundary = include;
108        self
109    }
110
111    /// Set whether to inject context
112    pub fn with_context_injection(mut self, inject: bool) -> Self {
113        self.inject_context = inject;
114        self
115    }
116
117    /// Add a section
118    pub fn add_section(&mut self, section: PromptSection) -> &mut Self {
119        self.sections.push(section);
120        self
121    }
122
123    /// Add multiple sections
124    pub fn add_sections(&mut self, sections: Vec<PromptSection>) -> &mut Self {
125        self.sections.extend(sections);
126        self
127    }
128
129    /// Clear all sections
130    pub fn clear_sections(&mut self) -> &mut Self {
131        self.sections.clear();
132        self
133    }
134
135    /// Invalidate cache (e.g., when context changes)
136    pub fn invalidate_cache(&mut self) {
137        self.cache.clear();
138        self.context_injector.invalidate();
139    }
140
141    /// Render a section with caching
142    fn render_section(&self, section: &PromptSection) -> String {
143        if section.cacheable {
144            let key = CacheKey::new(&section.name, self.profile.as_str());
145            self.cache.get_or_compute(&key, || section.compute_content())
146        } else {
147            section.compute_content()
148        }
149    }
150
151    /// Assemble the full prompt
152    pub fn assemble(&mut self) -> AssembledPrompt {
153        let mut cached_parts = Vec::new();
154        let mut dynamic_parts = Vec::new();
155        let mut cached_tokens = 0;
156        let mut dynamic_tokens = 0;
157
158        // Sort sections by order
159        let mut sections = self.sections.clone();
160        sections.sort_by_key(|s| s.order);
161
162        // Process each section
163        for section in &sections {
164            let content = self.render_section(section);
165            
166            if section.cacheable {
167                cached_parts.push((section.name.clone(), content.clone()));
168                cached_tokens += self.estimate_tokens(&content);
169            } else {
170                dynamic_parts.push((section.name.clone(), content.clone()));
171                dynamic_tokens += self.estimate_tokens(&content);
172            }
173        }
174
175        // Inject runtime context if enabled
176        let context_section = if self.inject_context {
177            let ctx = self.context_injector.render_full_context();
178            dynamic_tokens += self.estimate_tokens(&ctx);
179            Some(ctx)
180        } else {
181            None
182        };
183
184        // Assemble final prompt
185        let mut final_parts = Vec::new();
186
187        // Add cached parts first
188        for (name, content) in &cached_parts {
189            if !content.is_empty() {
190                final_parts.push(format!("[{}]\n{}", name, content));
191            }
192        }
193
194        // Add cache boundary if there are both cached and dynamic parts
195        if self.include_boundary && !cached_parts.is_empty() && (!dynamic_parts.is_empty() || context_section.is_some()) {
196            final_parts.push(CACHE_BOUNDARY.to_string());
197        }
198
199        // Add dynamic parts
200        for (name, content) in &dynamic_parts {
201            if !content.is_empty() {
202                final_parts.push(format!("[{}]\n{}", name, content));
203            }
204        }
205
206        // Add context section
207        if let Some(ctx) = context_section {
208            final_parts.push(ctx);
209        }
210
211        let full_prompt = final_parts.join("\n\n");
212
213        // Get cache stats
214        let stats = self.cache.stats();
215
216        AssembledPrompt {
217            prompt: full_prompt,
218            cached_sections: cached_parts.len(),
219            dynamic_sections: dynamic_parts.len(),
220            cached_tokens,
221            dynamic_tokens,
222            total_tokens: cached_tokens + dynamic_tokens,
223            cache_hit_rate: stats.hit_rate(),
224            profile: self.profile,
225        }
226    }
227
228    /// Assemble prompt for specific profile
229    pub fn assemble_for_profile(&mut self, profile: PromptProfile) -> AssembledPrompt {
230        self.profile = profile;
231        self.assemble()
232    }
233
234    /// Estimate token count
235    fn estimate_tokens(&self, content: &str) -> usize {
236        crate::prompt::cache::estimate_tokens(content)
237    }
238
239    /// Get cache statistics
240    pub fn cache_stats(&self) -> crate::prompt::cache::CacheStats {
241        self.cache.stats()
242    }
243
244    /// Get user context
245    pub fn get_user_context(&mut self) -> &UserContext {
246        self.context_injector.get_user_context()
247    }
248
249    /// Get system context
250    pub fn get_system_context(&mut self) -> &SystemContext {
251        self.context_injector.get_system_context()
252    }
253}
254
255/// Assembled prompt with metadata
256#[derive(Debug, Clone)]
257pub struct AssembledPrompt {
258    /// The full prompt text
259    pub prompt: String,
260    /// Number of cached sections
261    pub cached_sections: usize,
262    /// Number of dynamic sections
263    pub dynamic_sections: usize,
264    /// Estimated cached tokens
265    pub cached_tokens: usize,
266    /// Estimated dynamic tokens
267    pub dynamic_tokens: usize,
268    /// Total estimated tokens
269    pub total_tokens: usize,
270    /// Cache hit rate
271    pub cache_hit_rate: f64,
272    /// Profile used
273    pub profile: PromptProfile,
274}
275
276impl AssembledPrompt {
277    /// Check if prompt is empty
278    pub fn is_empty(&self) -> bool {
279        self.prompt.is_empty()
280    }
281
282    /// Get prompt length in characters
283    pub fn len(&self) -> usize {
284        self.prompt.len()
285    }
286
287    /// Get cache efficiency percentage
288    pub fn cache_efficiency(&self) -> f64 {
289        if self.total_tokens == 0 {
290            0.0
291        } else {
292            (self.cached_tokens as f64 / self.total_tokens as f64) * 100.0
293        }
294    }
295
296    /// Split prompt at cache boundary
297    pub fn split_at_boundary(&self) -> (Option<&str>, Option<&str>) {
298        if let Some(idx) = self.prompt.find(CACHE_BOUNDARY) {
299            let cached = &self.prompt[..idx];
300            let dynamic = &self.prompt[idx + CACHE_BOUNDARY.len()..];
301            (
302                if cached.is_empty() { None } else { Some(cached) },
303                if dynamic.is_empty() { None } else { Some(dynamic) },
304            )
305        } else {
306            if self.prompt.is_empty() {
307                (None, None)
308            } else {
309                (Some(&self.prompt), None)
310            }
311        }
312    }
313}
314
315/// Builder for creating prompt orchestrators
316pub struct PromptBuilder {
317    working_dir: std::path::PathBuf,
318    profile: PromptProfile,
319    sections: Vec<PromptSection>,
320    include_boundary: bool,
321    inject_context: bool,
322}
323
324impl PromptBuilder {
325    pub fn new<P: Into<std::path::PathBuf>>(working_dir: P) -> Self {
326        Self {
327            working_dir: working_dir.into(),
328            profile: PromptProfile::Default,
329            sections: Vec::new(),
330            include_boundary: true,
331            inject_context: true,
332        }
333    }
334
335    pub fn profile(mut self, profile: PromptProfile) -> Self {
336        self.profile = profile;
337        self
338    }
339
340    pub fn add_section(mut self, section: PromptSection) -> Self {
341        self.sections.push(section);
342        self
343    }
344
345    pub fn add_static(self, name: impl Into<String>, content: &'static str) -> Self {
346        self.add_section(PromptSection::static_section(name, content))
347    }
348
349    pub fn add_dynamic<F>(self, name: impl Into<String>, compute: F) -> Self
350    where
351        F: Fn() -> String + Send + Sync + 'static,
352    {
353        self.add_section(PromptSection::dynamic_section(name, compute))
354    }
355
356    pub fn no_boundary(mut self) -> Self {
357        self.include_boundary = false;
358        self
359    }
360
361    pub fn no_context(mut self) -> Self {
362        self.inject_context = false;
363        self
364    }
365
366    pub fn build(self) -> PromptOrchestrator {
367        let mut orchestrator = PromptOrchestrator::new(self.working_dir)
368            .with_profile(self.profile)
369            .with_boundary(self.include_boundary)
370            .with_context_injection(self.inject_context);
371        orchestrator.add_sections(self.sections);
372        orchestrator
373    }
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379
380    #[test]
381    fn test_assemble_simple() {
382        let mut orchestrator = PromptOrchestrator::new(std::env::current_dir().unwrap());
383        orchestrator.add_section(PromptSection::static_section("identity", "You are an AI assistant."));
384        
385        let assembled = orchestrator.assemble();
386        assert!(!assembled.prompt.is_empty());
387        assert!(assembled.prompt.contains("identity"));
388        assert!(assembled.cached_sections >= 1);
389    }
390
391    #[test]
392    fn test_assemble_with_dynamic() {
393        let mut orchestrator = PromptOrchestrator::new(std::env::current_dir().unwrap());
394        orchestrator.add_section(PromptSection::static_section("identity", "You are an AI."));
395        orchestrator.add_section(PromptSection::dynamic_section("date", || {
396            format!("Current date: {}", chrono::Local::now().format("%Y-%m-%d"))
397        }));
398        
399        let assembled = orchestrator.assemble();
400        assert!(assembled.dynamic_sections >= 1);
401        assert!(assembled.cached_sections >= 1);
402    }
403
404    #[test]
405    fn test_cache_boundary() {
406        let mut orchestrator = PromptOrchestrator::new(std::env::current_dir().unwrap())
407            .with_boundary(true)
408            .with_context_injection(false);
409        
410        orchestrator.add_section(PromptSection::static_section("cached", "cached content"));
411        orchestrator.add_section(PromptSection::dynamic_section("dynamic", || "dynamic content".to_string()));
412        
413        let assembled = orchestrator.assemble();
414        assert!(assembled.prompt.contains(CACHE_BOUNDARY));
415        
416        let (cached, dynamic) = assembled.split_at_boundary();
417        assert!(cached.is_some());
418        assert!(dynamic.is_some());
419    }
420
421    #[test]
422    fn test_profile() {
423        let orchestrator = PromptOrchestrator::new(std::env::current_dir().unwrap())
424            .with_profile(PromptProfile::Fast);
425        
426        assert_eq!(orchestrator.profile, PromptProfile::Fast);
427    }
428
429    #[test]
430    fn test_builder() {
431        let mut orchestrator = PromptBuilder::new(std::env::current_dir().unwrap())
432            .profile(PromptProfile::Review)
433            .add_static("identity", "You are a code reviewer.")
434            .add_dynamic("date", || "Today".to_string())
435            .build();
436        
437        let assembled = orchestrator.assemble();
438        assert!(assembled.prompt.contains("identity"));
439    }
440
441    #[test]
442    fn test_cache_efficiency() {
443        let mut orchestrator = PromptOrchestrator::new(std::env::current_dir().unwrap())
444            .with_context_injection(false);
445        
446        // Add static content with words (to have proper token estimate)
447        let static_content = "static content that should be cached properly test test test";
448        orchestrator.add_section(PromptSection::static_section("big", static_content));
449        orchestrator.add_section(PromptSection::dynamic_section("small", || "dynamic".to_string()));
450        
451        let assembled = orchestrator.assemble();
452        let efficiency = assembled.cache_efficiency();
453        
454        // Static content should be cached
455        assert!(efficiency >= 50.0, "Cache efficiency: {}", efficiency);
456    }
457
458    #[test]
459    fn test_invalidate_cache() {
460        let mut orchestrator = PromptOrchestrator::new(std::env::current_dir().unwrap());
461        orchestrator.add_section(PromptSection::static_section("test", "test content"));
462        
463        // First assembly
464        let _ = orchestrator.assemble();
465        
466        // Invalidate
467        orchestrator.invalidate_cache();
468        
469        // Should recalculate
470        let assembled = orchestrator.assemble();
471        assert!(assembled.prompt.contains("test"));
472    }
473}
matrixcode_core/prompt/orchestrator.rs

matrixcode_core/prompt/
orchestrator.rs
