cel_brief/sources/
user_message.rs

1//! [`UserMessageSource`] — pulls [`BriefContext::user_message`] into the brief.
2//!
3//! Critical-priority, never redactable. If `ctx.user_message` is `None` the
4//! source returns [`SourceError::Skipped`] — the builder treats that as zero
5//! contributions, not as a failure, so it's safe to register
6//! `UserMessageSource` unconditionally.
7
8use async_trait::async_trait;
9
10use crate::source::{Contribution, Source, SourceError};
11use crate::types::{BriefContext, Priority, Role, SourceId};
12
13/// A [`Source`] that mirrors [`BriefContext::user_message`] into the brief.
14///
15/// Skipped (not failed) when no user message is present. Critical priority so
16/// it survives any budget pruning short of `BudgetUnsatisfiable`. Emits a
17/// [`Role::User`] text contribution; not redactable — the user's literal
18/// words should reach the model verbatim.
19#[derive(Debug, Clone)]
20pub struct UserMessageSource {
21    id: SourceId,
22}
23
24impl UserMessageSource {
25    /// Construct with the default ID `"user_message"`.
26    pub fn new() -> Self {
27        UserMessageSource {
28            id: SourceId::new("user_message"),
29        }
30    }
31
32    /// Override the default [`SourceId`].
33    pub fn with_id(mut self, id: impl Into<SourceId>) -> Self {
34        self.id = id.into();
35        self
36    }
37}
38
39impl Default for UserMessageSource {
40    fn default() -> Self {
41        UserMessageSource::new()
42    }
43}
44
45#[async_trait]
46impl Source for UserMessageSource {
47    fn id(&self) -> SourceId {
48        self.id.clone()
49    }
50
51    fn priority(&self) -> Priority {
52        Priority::Critical
53    }
54
55    async fn contribute(&self, ctx: &BriefContext) -> Result<Vec<Contribution>, SourceError> {
56        let Some(msg) = ctx.user_message.as_deref() else {
57            return Err(SourceError::Skipped(
58                "BriefContext.user_message is None".into(),
59            ));
60        };
61        let est = msg.len().div_ceil(4);
62        Ok(vec![Contribution::text(Role::User, msg.to_owned(), est)
63            .with_importance(1.0)
64            .with_redactable(false)])
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use crate::source::ContributionContent;
72    use crate::types::TokenBudget;
73
74    #[tokio::test]
75    async fn round_trips_user_message() {
76        let src = UserMessageSource::new();
77        let ctx = BriefContext::new(TokenBudget::default()).with_user_message("hi there");
78        let contributions = src.contribute(&ctx).await.expect("contribute ok");
79        assert_eq!(contributions.len(), 1);
80        let c = &contributions[0];
81        assert_eq!(c.importance, 1.0);
82        assert!(!c.redactable);
83        match &c.content {
84            ContributionContent::Text { role, content } => {
85                assert_eq!(*role, Role::User);
86                assert_eq!(content, "hi there");
87            }
88            other => panic!("expected Text, got {other:?}"),
89        }
90    }
91
92    #[tokio::test]
93    async fn skipped_when_no_user_message() {
94        let src = UserMessageSource::default();
95        let ctx = BriefContext::new(TokenBudget::default());
96        let err = src.contribute(&ctx).await.expect_err("should skip");
97        match err {
98            SourceError::Skipped(_) => {}
99            other => panic!("expected Skipped, got {other:?}"),
100        }
101    }
102
103    #[tokio::test]
104    async fn priority_is_critical() {
105        assert_eq!(UserMessageSource::new().priority(), Priority::Critical);
106    }
107
108    #[tokio::test]
109    async fn with_id_overrides_default() {
110        let src = UserMessageSource::new().with_id("user_say");
111        assert_eq!(src.id(), SourceId::new("user_say"));
112    }
113}
cel_brief/sources/user_message.rs

cel_brief/sources/
user_message.rs
