cel_brief/sources/
system_prompt.rs1use async_trait::async_trait;
11
12use crate::source::{Contribution, Source, SourceError};
13use crate::types::{BriefContext, Priority, SourceId};
14
15#[derive(Debug, Clone)]
22pub struct SystemPromptSource {
23 id: SourceId,
24 text: String,
25 estimated_tokens: usize,
26}
27
28impl SystemPromptSource {
29 pub fn new(text: impl Into<String>) -> Self {
31 let text = text.into();
32 let estimated_tokens = text.len().div_ceil(4);
33 SystemPromptSource {
34 id: SourceId::new("system_prompt"),
35 text,
36 estimated_tokens,
37 }
38 }
39
40 pub fn with_id(mut self, id: impl Into<SourceId>) -> Self {
44 self.id = id.into();
45 self
46 }
47
48 pub fn with_estimated_tokens(mut self, tokens: usize) -> Self {
51 self.estimated_tokens = tokens;
52 self
53 }
54
55 pub fn text(&self) -> &str {
57 &self.text
58 }
59}
60
61#[async_trait]
62impl Source for SystemPromptSource {
63 fn id(&self) -> SourceId {
64 self.id.clone()
65 }
66
67 fn priority(&self) -> Priority {
68 Priority::Critical
69 }
70
71 async fn contribute(&self, _ctx: &BriefContext) -> Result<Vec<Contribution>, SourceError> {
72 Ok(vec![Contribution::system(
73 self.text.clone(),
74 self.estimated_tokens,
75 )])
76 }
77}
78
79#[cfg(test)]
80mod tests {
81 use super::*;
82 use crate::source::ContributionContent;
83 use crate::types::TokenBudget;
84
85 #[tokio::test]
86 async fn emits_one_critical_system_contribution() {
87 let src = SystemPromptSource::new("You are a helpful assistant.");
88 assert_eq!(src.priority(), Priority::Critical);
89 assert_eq!(src.id(), SourceId::new("system_prompt"));
90
91 let ctx = BriefContext::new(TokenBudget::default());
92 let contributions = src.contribute(&ctx).await.expect("contribute ok");
93 assert_eq!(contributions.len(), 1);
94 let c = &contributions[0];
95 assert!(!c.redactable);
96 assert_eq!(c.importance, 1.0);
97 match &c.content {
98 ContributionContent::System { text } => {
99 assert_eq!(text, "You are a helpful assistant.");
100 }
101 other => panic!("expected System, got {other:?}"),
102 }
103 }
104
105 #[tokio::test]
106 async fn with_id_overrides_default() {
107 let src = SystemPromptSource::new("hi").with_id("personality");
108 assert_eq!(src.id(), SourceId::new("personality"));
109 }
110
111 #[tokio::test]
112 async fn with_estimated_tokens_overrides_default() {
113 let src = SystemPromptSource::new("hi").with_estimated_tokens(99);
114 let ctx = BriefContext::new(TokenBudget::default());
115 let c = src.contribute(&ctx).await.expect("contribute ok");
116 assert_eq!(c[0].estimated_tokens, 99);
117 }
118
119 #[test]
120 fn default_estimate_is_four_chars_per_token() {
121 let src = SystemPromptSource::new("hello, world");
123 assert_eq!(src.estimated_tokens, 3);
124 }
125}