walrus_core/model/
message.rs

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

walrus_core/model/
message.rs
