rustlavel_ai/
message.rs

1//! The conversation model, kept deliberately free of any provider's shape.
2//!
3//! Every provider is translated to and from these types, so an application
4//! that switches from Anthropic to Ollama changes one line of configuration
5//! and nothing else.
6
7use rustlavel_core::Json;
8
9/// Who a message came from.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum Role {
12    /// Instructions that frame the whole conversation.
13    System,
14    User,
15    Assistant,
16    /// The result of a tool the assistant asked to run.
17    Tool,
18}
19
20impl Role {
21    pub fn as_str(self) -> &'static str {
22        match self {
23            Role::System => "system",
24            Role::User => "user",
25            Role::Assistant => "assistant",
26            Role::Tool => "tool",
27        }
28    }
29
30    pub fn parse(text: &str) -> Option<Role> {
31        match text {
32            "system" => Some(Role::System),
33            "user" => Some(Role::User),
34            "assistant" => Some(Role::Assistant),
35            "tool" => Some(Role::Tool),
36            _ => None,
37        }
38    }
39}
40
41impl std::fmt::Display for Role {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        f.write_str(self.as_str())
44    }
45}
46
47/// One part of a message.
48///
49/// A message is a list of parts rather than a string because a single
50/// assistant turn can say something *and* ask for two tools, and every
51/// provider needs those parts kept apart to replay the turn back to it.
52#[derive(Debug, Clone, PartialEq)]
53pub enum Content {
54    Text(String),
55    /// The assistant asking for a tool to be run.
56    ToolUse { id: String, name: String, input: Json },
57    /// What the tool answered, fed back on the next turn.
58    ToolResult { id: String, name: String, output: Json, is_error: bool },
59}
60
61impl Content {
62    pub fn text(text: impl Into<String>) -> Content {
63        Content::Text(text.into())
64    }
65
66    /// The text of this part, or `None` for a tool part.
67    pub fn as_text(&self) -> Option<&str> {
68        match self {
69            Content::Text(text) => Some(text),
70            _ => None,
71        }
72    }
73}
74
75/// One turn in a conversation.
76#[derive(Debug, Clone, PartialEq)]
77pub struct Message {
78    pub role: Role,
79    pub content: Vec<Content>,
80}
81
82impl Message {
83    pub fn new(role: Role, content: Vec<Content>) -> Message {
84        Message { role, content }
85    }
86
87    pub fn system(text: impl Into<String>) -> Message {
88        Message::new(Role::System, vec![Content::text(text)])
89    }
90
91    pub fn user(text: impl Into<String>) -> Message {
92        Message::new(Role::User, vec![Content::text(text)])
93    }
94
95    pub fn assistant(text: impl Into<String>) -> Message {
96        Message::new(Role::Assistant, vec![Content::text(text)])
97    }
98
99    /// The answer from a tool the assistant asked for.
100    pub fn tool_result(id: impl Into<String>, name: impl Into<String>, output: Json) -> Message {
101        Message::new(
102            Role::Tool,
103            vec![Content::ToolResult {
104                id: id.into(),
105                name: name.into(),
106                output,
107                is_error: false,
108            }],
109        )
110    }
111
112    /// A tool that failed.
113    ///
114    /// Handed back to the model rather than raised, because a model that hears
115    /// "that city does not exist" usually corrects itself, and a model that
116    /// hears nothing at all just hangs the loop.
117    pub fn tool_error(
118        id: impl Into<String>,
119        name: impl Into<String>,
120        message: impl Into<String>,
121    ) -> Message {
122        Message::new(
123            Role::Tool,
124            vec![Content::ToolResult {
125                id: id.into(),
126                name: name.into(),
127                output: Json::from(message.into()),
128                is_error: true,
129            }],
130        )
131    }
132
133    pub fn with(mut self, part: Content) -> Message {
134        self.content.push(part);
135        self
136    }
137
138    /// Every text part joined, which is what a caller usually means by
139    /// "what did it say".
140    pub fn text(&self) -> String {
141        self.content
142            .iter()
143            .filter_map(Content::as_text)
144            .collect::<Vec<_>>()
145            .join("")
146    }
147
148    /// The tool requests in this turn, if any.
149    pub fn tool_uses(&self) -> impl Iterator<Item = (&str, &str, &Json)> {
150        self.content.iter().filter_map(|part| match part {
151            Content::ToolUse { id, name, input } => Some((id.as_str(), name.as_str(), input)),
152            _ => None,
153        })
154    }
155
156    pub fn is_empty(&self) -> bool {
157        self.content.is_empty()
158    }
159}
160
161/// An ordered list of messages, the thing you actually send.
162///
163/// Thin on purpose: the value is in having one type both the front door and
164/// the tool loop can append to without arguing about ownership.
165#[derive(Debug, Clone, Default, PartialEq)]
166pub struct Conversation {
167    messages: Vec<Message>,
168}
169
170impl Conversation {
171    pub fn new() -> Conversation {
172        Conversation::default()
173    }
174
175    pub fn system(mut self, text: impl Into<String>) -> Conversation {
176        self.messages.push(Message::system(text));
177        self
178    }
179
180    pub fn user(mut self, text: impl Into<String>) -> Conversation {
181        self.messages.push(Message::user(text));
182        self
183    }
184
185    pub fn assistant(mut self, text: impl Into<String>) -> Conversation {
186        self.messages.push(Message::assistant(text));
187        self
188    }
189
190    pub fn push(&mut self, message: Message) {
191        self.messages.push(message);
192    }
193
194    pub fn extend(&mut self, messages: impl IntoIterator<Item = Message>) {
195        self.messages.extend(messages);
196    }
197
198    pub fn messages(&self) -> &[Message] {
199        &self.messages
200    }
201
202    pub fn last(&self) -> Option<&Message> {
203        self.messages.last()
204    }
205
206    pub fn len(&self) -> usize {
207        self.messages.len()
208    }
209
210    pub fn is_empty(&self) -> bool {
211        self.messages.is_empty()
212    }
213
214    pub fn into_messages(self) -> Vec<Message> {
215        self.messages
216    }
217}
218
219impl From<Conversation> for Vec<Message> {
220    fn from(conversation: Conversation) -> Vec<Message> {
221        conversation.messages
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    #[test]
230    fn a_message_joins_its_text_parts_and_keeps_tool_parts_apart() {
231        let message = Message::assistant("Let me check. ")
232            .with(Content::text("One moment."))
233            .with(Content::ToolUse {
234                id: "toolu_1".into(),
235                name: "get_weather".into(),
236                input: Json::object([("city", "Oslo".into())]),
237            });
238
239        assert_eq!(message.text(), "Let me check. One moment.");
240
241        let uses: Vec<_> = message.tool_uses().collect();
242        assert_eq!(uses.len(), 1);
243        assert_eq!(uses[0].0, "toolu_1");
244        assert_eq!(uses[0].1, "get_weather");
245        assert_eq!(uses[0].2.get("city").unwrap().as_str(), Some("Oslo"));
246    }
247
248    #[test]
249    fn a_failing_tool_becomes_a_message_rather_than_an_error() {
250        let message = Message::tool_error("toolu_1", "get_weather", "no such city");
251
252        assert_eq!(message.role, Role::Tool);
253        match &message.content[0] {
254            Content::ToolResult { is_error, output, .. } => {
255                assert!(is_error);
256                assert_eq!(output.as_str(), Some("no such city"));
257            }
258            other => panic!("expected a tool result, got {other:?}"),
259        }
260    }
261
262    #[test]
263    fn a_conversation_builds_up_in_order() {
264        let mut conversation = Conversation::new().system("Be brief.").user("Hello");
265        conversation.push(Message::assistant("Hi."));
266
267        assert_eq!(conversation.len(), 3);
268        assert_eq!(conversation.messages()[0].role, Role::System);
269        assert_eq!(conversation.last().unwrap().text(), "Hi.");
270        assert!(!conversation.is_empty());
271    }
272
273    #[test]
274    fn roles_round_trip_through_their_wire_names() {
275        for role in [Role::System, Role::User, Role::Assistant, Role::Tool] {
276            assert_eq!(Role::parse(role.as_str()), Some(role));
277            assert_eq!(role.to_string(), role.as_str());
278        }
279        assert_eq!(Role::parse("wizard"), None);
280    }
281}
rustlavel_ai/message.rs

rustlavel_ai/
message.rs
