crabtalk_core/model/
message.rs

1//! Turbofish LLM message
2
3use crate::model::{StreamChunk, ToolCall};
4pub use crabllm_core::Role;
5use serde::{Deserialize, Serialize};
6use std::collections::BTreeMap;
7
8/// A message in the chat
9#[derive(Debug, Clone, Deserialize, Serialize)]
10pub struct Message {
11    /// The role of the message
12    pub role: Role,
13
14    /// The content of the message
15    #[serde(default, skip_serializing_if = "String::is_empty")]
16    pub content: String,
17
18    /// The reasoning content
19    #[serde(default, skip_serializing_if = "String::is_empty")]
20    pub reasoning_content: String,
21
22    /// The tool call id
23    #[serde(default, skip_serializing_if = "String::is_empty")]
24    pub tool_call_id: String,
25
26    /// The tool calls
27    #[serde(default, skip_serializing_if = "Vec::is_empty")]
28    pub tool_calls: Vec<ToolCall>,
29
30    /// The sender identity (runtime-only, never serialized to providers).
31    ///
32    /// Convention: empty = local/owner, `"tg:12345"` = Telegram user.
33    #[serde(skip)]
34    pub sender: String,
35
36    /// Whether this message was auto-injected by the runtime (e.g. recall).
37    /// Auto-injected messages are stripped before each new run.
38    #[serde(skip)]
39    pub auto_injected: bool,
40}
41
42impl Message {
43    /// Create a new system message
44    pub fn system(content: impl Into<String>) -> Self {
45        Self {
46            role: Role::System,
47            content: content.into(),
48            ..Default::default()
49        }
50    }
51
52    /// Create a new user message
53    pub fn user(content: impl Into<String>) -> Self {
54        Self {
55            role: Role::User,
56            content: content.into(),
57            ..Default::default()
58        }
59    }
60
61    /// Create a new user message with sender identity.
62    pub fn user_with_sender(content: impl Into<String>, sender: impl Into<String>) -> Self {
63        Self {
64            role: Role::User,
65            content: content.into(),
66            sender: sender.into(),
67            ..Default::default()
68        }
69    }
70
71    /// Create a new assistant message
72    pub fn assistant(
73        content: impl Into<String>,
74        reasoning: Option<String>,
75        tool_calls: Option<&[ToolCall]>,
76    ) -> Self {
77        Self {
78            role: Role::Assistant,
79            content: content.into(),
80            reasoning_content: reasoning.unwrap_or_default(),
81            tool_calls: tool_calls.map(|tc| tc.to_vec()).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<String>) -> 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    /// Peek at accumulated tool calls with non-empty names.
165    /// Returns clones of the current state (args may be partial).
166    pub fn peek_tool_calls(&self) -> Vec<ToolCall> {
167        self.calls
168            .values()
169            .filter(|c| !c.function.name.is_empty())
170            .cloned()
171            .collect()
172    }
173
174    /// Build the message
175    pub fn build(mut self) -> Message {
176        if !self.calls.is_empty() {
177            self.message.tool_calls = self.calls.into_values().collect();
178        }
179        self.message
180    }
181}
182
183impl Default for Message {
184    fn default() -> Self {
185        Self {
186            role: Role::User,
187            content: String::new(),
188            reasoning_content: String::new(),
189            tool_call_id: String::new(),
190            tool_calls: Vec::new(),
191            sender: String::new(),
192            auto_injected: false,
193        }
194    }
195}
crabtalk_core/model/message.rs

crabtalk_core/model/
message.rs
