Skip to main content

use_ai_message/
lib.rs

1#![forbid(unsafe_code)]
2#![doc = include_str!("../README.md")]
3
4use core::{fmt, str::FromStr};
5use std::error::Error;
6
7pub mod prelude {
8    pub use crate::{
9        AiConversationId, AiMessageContentRef, AiMessageError, AiMessageFormat, AiMessageId,
10        AiMessagePartKind, AiMessageRole, AiMessageSource, AiMessageStatus, AiMessageTokenCount,
11        AiMessageVisibility,
12    };
13}
14
15macro_rules! message_text_newtype {
16    ($name:ident) => {
17        #[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
18        pub struct $name(String);
19
20        impl $name {
21            pub fn new(value: impl AsRef<str>) -> Result<Self, AiMessageError> {
22                non_empty_text(value).map(Self)
23            }
24
25            pub fn as_str(&self) -> &str {
26                &self.0
27            }
28
29            pub fn value(&self) -> &str {
30                self.as_str()
31            }
32
33            pub fn into_string(self) -> String {
34                self.0
35            }
36        }
37
38        impl AsRef<str> for $name {
39            fn as_ref(&self) -> &str {
40                self.as_str()
41            }
42        }
43
44        impl fmt::Display for $name {
45            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
46                formatter.write_str(self.as_str())
47            }
48        }
49
50        impl FromStr for $name {
51            type Err = AiMessageError;
52
53            fn from_str(value: &str) -> Result<Self, Self::Err> {
54                Self::new(value)
55            }
56        }
57
58        impl TryFrom<&str> for $name {
59            type Error = AiMessageError;
60
61            fn try_from(value: &str) -> Result<Self, Self::Error> {
62                Self::new(value)
63            }
64        }
65    };
66}
67
68macro_rules! message_enum {
69    ($name:ident { $($variant:ident => $label:literal),+ $(,)? }) => {
70        #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
71        pub enum $name {
72            $($variant),+
73        }
74
75        impl $name {
76            pub const ALL: &'static [Self] = &[$(Self::$variant),+];
77
78            pub const fn as_str(self) -> &'static str {
79                match self {
80                    $(Self::$variant => $label),+
81                }
82            }
83        }
84
85        impl fmt::Display for $name {
86            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
87                formatter.write_str(self.as_str())
88            }
89        }
90
91        impl FromStr for $name {
92            type Err = AiMessageError;
93
94            fn from_str(value: &str) -> Result<Self, Self::Err> {
95                match normalized_label(value)?.as_str() {
96                    $($label => Ok(Self::$variant),)+
97                    _ => Err(AiMessageError::UnknownLabel),
98                }
99            }
100        }
101    };
102}
103
104message_text_newtype!(AiMessageId);
105message_text_newtype!(AiConversationId);
106message_text_newtype!(AiMessageContentRef);
107
108#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
109pub struct AiMessageTokenCount(u32);
110
111impl AiMessageTokenCount {
112    pub const fn new(value: u32) -> Self {
113        Self(value)
114    }
115
116    pub const fn value(self) -> u32 {
117        self.0
118    }
119
120    pub const fn get(self) -> u32 {
121        self.0
122    }
123}
124
125impl fmt::Display for AiMessageTokenCount {
126    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
127        self.0.fmt(formatter)
128    }
129}
130
131message_enum!(AiMessageRole {
132    System => "system",
133    Developer => "developer",
134    User => "user",
135    Assistant => "assistant",
136    Tool => "tool",
137    Function => "function",
138    Observer => "observer",
139    Evaluator => "evaluator",
140    Unknown => "unknown",
141});
142
143message_enum!(AiMessagePartKind {
144    Text => "text",
145    Image => "image",
146    Audio => "audio",
147    Video => "video",
148    File => "file",
149    ToolCall => "tool-call",
150    ToolResult => "tool-result",
151    Citation => "citation",
152    ThoughtSummary => "thought-summary",
153    Metadata => "metadata",
154    Custom => "custom",
155});
156
157message_enum!(AiMessageStatus {
158    Draft => "draft",
159    Pending => "pending",
160    Streaming => "streaming",
161    Complete => "complete",
162    Failed => "failed",
163    Cancelled => "cancelled",
164    Redacted => "redacted",
165});
166
167message_enum!(AiMessageFormat {
168    PlainText => "plain-text",
169    Markdown => "markdown",
170    Json => "json",
171    Xml => "xml",
172    Yaml => "yaml",
173    Multipart => "multipart",
174    Custom => "custom",
175});
176
177message_enum!(AiMessageVisibility {
178    Internal => "internal",
179    UserVisible => "user-visible",
180    Hidden => "hidden",
181    Redacted => "redacted",
182});
183
184message_enum!(AiMessageSource {
185    Human => "human",
186    Model => "model",
187    Tool => "tool",
188    System => "system",
189    Import => "import",
190    Synthetic => "synthetic",
191    Unknown => "unknown",
192});
193
194#[derive(Clone, Copy, Debug, Eq, PartialEq)]
195pub enum AiMessageError {
196    Empty,
197    UnknownLabel,
198}
199
200impl fmt::Display for AiMessageError {
201    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
202        match self {
203            Self::Empty => formatter.write_str("AI message metadata text cannot be empty"),
204            Self::UnknownLabel => formatter.write_str("unknown AI message metadata label"),
205        }
206    }
207}
208
209impl Error for AiMessageError {}
210
211fn non_empty_text(value: impl AsRef<str>) -> Result<String, AiMessageError> {
212    let trimmed = value.as_ref().trim();
213    if trimmed.is_empty() {
214        Err(AiMessageError::Empty)
215    } else {
216        Ok(trimmed.to_string())
217    }
218}
219
220fn normalized_label(value: &str) -> Result<String, AiMessageError> {
221    let trimmed = value.trim();
222    if trimmed.is_empty() {
223        Err(AiMessageError::Empty)
224    } else {
225        Ok(trimmed.to_ascii_lowercase().replace(['_', ' '], "-"))
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::{
232        AiConversationId, AiMessageContentRef, AiMessageError, AiMessageFormat, AiMessageId,
233        AiMessagePartKind, AiMessageRole, AiMessageSource, AiMessageStatus, AiMessageTokenCount,
234        AiMessageVisibility,
235    };
236    use core::{fmt, str::FromStr};
237
238    macro_rules! assert_text_newtype {
239        ($type:ty, $value:literal) => {{
240            let value = <$type>::new(concat!(" ", $value, " "))?;
241            assert_eq!(value.as_str(), $value);
242            assert_eq!(value.value(), $value);
243            assert_eq!(value.as_ref(), $value);
244            assert_eq!(value.to_string(), $value);
245            assert_eq!(<$type as TryFrom<&str>>::try_from($value)?, value);
246            assert_eq!(value.into_string(), $value.to_string());
247        }};
248    }
249
250    fn assert_enum_family<T>(variants: &[T]) -> Result<(), AiMessageError>
251    where
252        T: Copy + Eq + fmt::Debug + fmt::Display + FromStr<Err = AiMessageError>,
253    {
254        for variant in variants {
255            let label = variant.to_string();
256            assert_eq!(label.parse::<T>()?, *variant);
257            assert_eq!(label.replace('-', "_").parse::<T>()?, *variant);
258            assert_eq!(label.replace('-', " ").parse::<T>()?, *variant);
259        }
260        Ok(())
261    }
262
263    #[test]
264    fn validates_message_text_newtypes() -> Result<(), AiMessageError> {
265        assert_text_newtype!(AiMessageId, "msg-001");
266        assert_text_newtype!(AiConversationId, "conv-001");
267        assert_text_newtype!(AiMessageContentRef, "content:001");
268        assert_eq!(AiMessageId::new("  "), Err(AiMessageError::Empty));
269        Ok(())
270    }
271
272    #[test]
273    fn validates_token_counts() {
274        let zero = AiMessageTokenCount::new(0);
275        let count = AiMessageTokenCount::new(42);
276
277        assert_eq!(zero.value(), 0);
278        assert_eq!(count.get(), 42);
279        assert_eq!(count.to_string(), "42");
280    }
281
282    #[test]
283    fn displays_and_parses_message_enums() -> Result<(), AiMessageError> {
284        assert_enum_family(AiMessageRole::ALL)?;
285        assert_enum_family(AiMessagePartKind::ALL)?;
286        assert_enum_family(AiMessageStatus::ALL)?;
287        assert_enum_family(AiMessageFormat::ALL)?;
288        assert_enum_family(AiMessageVisibility::ALL)?;
289        assert_enum_family(AiMessageSource::ALL)?;
290        assert_eq!(
291            "user visible".parse::<AiMessageVisibility>()?,
292            AiMessageVisibility::UserVisible
293        );
294        Ok(())
295    }
296}