walrus_core/model/
message.rs

1//! Turbofish LLM message
2
3use crate::model::{StreamChunk, ToolCall};
4use compact_str::CompactString;
5use serde::{Deserialize, Serialize};
6use smallvec::SmallVec;
7use std::collections::BTreeMap;
8
9/// A message in the chat
10#[derive(Debug, Clone, Deserialize, Serialize, Default)]
11pub struct Message {
12    /// The role of the message
13    pub role: Role,
14
15    /// The content of the message
16    #[serde(skip_serializing_if = "String::is_empty")]
17    pub content: String,
18
19    /// The reasoning content
20    #[serde(skip_serializing_if = "String::is_empty")]
21    pub reasoning_content: String,
22
23    /// The tool call id
24    #[serde(skip_serializing_if = "CompactString::is_empty")]
25    pub tool_call_id: CompactString,
26
27    /// The tool calls
28    #[serde(skip_serializing_if = "SmallVec::is_empty")]
29    pub tool_calls: SmallVec<[ToolCall; 4]>,
30
31    /// The sender identity (runtime-only, never serialized to providers).
32    ///
33    /// Convention: empty = local/owner, `"tg:12345"` = Telegram user,
34    /// `"dc:67890"` = Discord user.
35    #[serde(skip)]
36    pub sender: CompactString,
37}
38
39impl Message {
40    /// Create a new system message
41    pub fn system(content: impl Into<String>) -> Self {
42        Self {
43            role: Role::System,
44            content: content.into(),
45            ..Default::default()
46        }
47    }
48
49    /// Create a new user message
50    pub fn user(content: impl Into<String>) -> Self {
51        Self {
52            role: Role::User,
53            content: content.into(),
54            ..Default::default()
55        }
56    }
57
58    /// Create a new user message with sender identity.
59    pub fn user_with_sender(content: impl Into<String>, sender: impl Into<CompactString>) -> Self {
60        Self {
61            role: Role::User,
62            content: content.into(),
63            sender: sender.into(),
64            ..Default::default()
65        }
66    }
67
68    /// Create a new assistant message
69    pub fn assistant(
70        content: impl Into<String>,
71        reasoning: Option<String>,
72        tool_calls: Option<&[ToolCall]>,
73    ) -> Self {
74        Self {
75            role: Role::Assistant,
76            content: content.into(),
77            reasoning_content: reasoning.unwrap_or_default(),
78            tool_calls: tool_calls
79                .map(|tc| tc.iter().cloned().collect())
80                .unwrap_or_default(),
81            ..Default::default()
82        }
83    }
84
85    /// Create a new tool message
86    pub fn tool(content: impl Into<String>, call: impl Into<CompactString>) -> Self {
87        Self {
88            role: Role::Tool,
89            content: content.into(),
90            tool_call_id: call.into(),
91            ..Default::default()
92        }
93    }
94
95    /// Create a new message builder
96    pub fn builder(role: Role) -> MessageBuilder {
97        MessageBuilder::new(role)
98    }
99
100    /// Estimate the number of tokens in this message.
101    ///
102    /// Uses a simple heuristic: ~4 characters per token.
103    pub fn estimate_tokens(&self) -> usize {
104        let chars = self.content.len()
105            + self.reasoning_content.len()
106            + self.tool_call_id.len()
107            + self
108                .tool_calls
109                .iter()
110                .map(|tc| tc.function.name.len() + tc.function.arguments.len())
111                .sum::<usize>();
112        (chars / 4).max(1)
113    }
114}
115
116/// Estimate total tokens across a slice of messages.
117pub fn estimate_tokens(messages: &[Message]) -> usize {
118    messages.iter().map(|m| m.estimate_tokens()).sum()
119}
120
121/// A builder for messages
122pub struct MessageBuilder {
123    /// The message
124    message: Message,
125    /// The tool calls
126    calls: BTreeMap<u32, ToolCall>,
127}
128
129impl MessageBuilder {
130    /// Create a new message builder
131    pub fn new(role: Role) -> Self {
132        Self {
133            message: Message {
134                role,
135                ..Default::default()
136            },
137            calls: BTreeMap::new(),
138        }
139    }
140
141    /// Accept a chunk from the stream
142    pub fn accept(&mut self, chunk: &StreamChunk) -> bool {
143        if let Some(calls) = chunk.tool_calls() {
144            for call in calls {
145                let entry = self.calls.entry(call.index).or_default();
146                entry.merge(call);
147            }
148        }
149
150        let mut has_content = false;
151        if let Some(content) = chunk.content() {
152            self.message.content.push_str(content);
153            has_content = true;
154        }
155
156        if let Some(reason) = chunk.reasoning_content() {
157            self.message.reasoning_content.push_str(reason);
158        }
159
160        has_content
161    }
162
163    /// Build the message
164    pub fn build(mut self) -> Message {
165        if !self.calls.is_empty() {
166            self.message.tool_calls = self.calls.into_values().collect();
167        }
168        self.message
169    }
170}
171
172/// The role of a message
173#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Deserialize, Serialize, Default)]
174pub enum Role {
175    /// The user role
176    #[serde(rename = "user")]
177    #[default]
178    User,
179    /// The assistant role
180    #[serde(rename = "assistant")]
181    Assistant,
182    /// The system role
183    #[serde(rename = "system")]
184    System,
185    /// The tool role
186    #[serde(rename = "tool")]
187    Tool,
188}
walrus_core/model/message.rs

walrus_core/model/
message.rs
