1use serde::{Deserialize, Serialize};
2use serde_json::{json, Value};
3
4#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
5pub struct ChatMessage {
6 pub role: String,
7 #[serde(default, skip_serializing_if = "Option::is_none")]
8 pub content: Option<String>,
9 #[serde(default, skip_serializing_if = "Option::is_none")]
10 pub reasoning_details: Option<Vec<Value>>,
11 #[serde(default, skip_serializing_if = "Option::is_none")]
12 pub name: Option<String>,
13 #[serde(default, skip_serializing_if = "Option::is_none")]
14 pub tool_call_id: Option<String>,
15 #[serde(default, skip_serializing_if = "Vec::is_empty")]
16 pub tool_calls: Vec<ChatToolCall>,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
20pub struct ChatToolCall {
21 pub id: String,
22 pub name: String,
23 pub arguments: String,
24}
25
26pub(crate) fn estimate_message_tokens(message: &ChatMessage) -> usize {
34 serde_json::to_vec(message)
35 .map(|encoded| encoded.len().div_ceil(4).max(1))
36 .unwrap_or(1)
37}
38
39pub(crate) fn estimate_context_tokens(messages: &[ChatMessage]) -> usize {
40 messages
41 .iter()
42 .map(estimate_message_tokens)
43 .sum::<usize>()
44 .max(1)
45}
46
47impl ChatMessage {
48 pub fn system(content: String) -> Self {
49 Self {
50 role: "system".to_owned(),
51 content: Some(content),
52 reasoning_details: None,
53 name: None,
54 tool_call_id: None,
55 tool_calls: Vec::new(),
56 }
57 }
58
59 pub fn user(content: String) -> Self {
60 Self {
61 role: "user".to_owned(),
62 content: Some(content),
63 reasoning_details: None,
64 name: None,
65 tool_call_id: None,
66 tool_calls: Vec::new(),
67 }
68 }
69
70 pub fn assistant(content: String, tool_calls: Vec<ChatToolCall>) -> Self {
71 Self {
72 role: "assistant".to_owned(),
73 content: (!content.is_empty()).then_some(content),
74 reasoning_details: None,
75 name: None,
76 tool_call_id: None,
77 tool_calls,
78 }
79 }
80
81 pub fn tool(tool_call_id: String, name: String, content: String) -> Self {
82 Self {
83 role: "tool".to_owned(),
84 content: Some(content),
85 reasoning_details: None,
86 name: Some(name),
87 tool_call_id: Some(tool_call_id),
88 tool_calls: Vec::new(),
89 }
90 }
91
92 pub fn to_openai_value(&self) -> Value {
93 let mut message = json!({
94 "role": self.role,
95 "content": self.content,
96 });
97 if self.role == "assistant" {
98 if let Some(reasoning_details) = &self.reasoning_details {
99 message["reasoning_details"] = Value::Array(reasoning_details.clone());
100 }
101 }
102 if let Some(name) = &self.name {
103 message["name"] = Value::String(name.clone());
104 }
105 if let Some(tool_call_id) = &self.tool_call_id {
106 message["tool_call_id"] = Value::String(tool_call_id.clone());
107 }
108 if !self.tool_calls.is_empty() {
109 message["tool_calls"] = Value::Array(
110 self.tool_calls
111 .iter()
112 .map(|tool_call| {
113 json!({
114 "id": tool_call.id,
115 "type": "function",
116 "function": {
117 "name": tool_call.name,
118 "arguments": tool_call.arguments,
119 }
120 })
121 })
122 .collect(),
123 );
124 }
125 message
126 }
127}
128
129#[cfg(test)]
130mod tests {
131 use super::*;
132
133 #[test]
134 fn context_token_estimate_is_nonzero_and_grows_with_messages() {
135 let empty = estimate_context_tokens(&[]);
136 let one = estimate_context_tokens(&[ChatMessage::user("hello".to_owned())]);
137
138 assert_eq!(empty, 1);
139 assert!(one > empty);
140 assert!(estimate_message_tokens(&ChatMessage::user("hello".to_owned())) > 0);
141 }
142
143 #[test]
144 fn tool_assistant_messages_have_openai_compatible_shape() {
145 let assistant = ChatMessage::assistant(
146 String::new(),
147 vec![ChatToolCall {
148 id: "call-1".to_owned(),
149 name: "cmd".to_owned(),
150 arguments: r#"{"command":"pwd"}"#.to_owned(),
151 }],
152 );
153 assert_eq!(assistant.to_openai_value()["content"], Value::Null);
154 assert_eq!(
155 assistant.to_openai_value()["tool_calls"][0]["type"],
156 "function"
157 );
158
159 let tool = ChatMessage::tool(
160 "call-1".to_owned(),
161 "cmd".to_owned(),
162 "{\"exit_code\":0}".to_owned(),
163 );
164 assert_eq!(tool.to_openai_value()["tool_call_id"], "call-1");
165 assert_eq!(tool.to_openai_value()["name"], "cmd");
166 }
167
168 #[test]
169 fn reasoning_details_are_optional_and_only_sent_for_assistant_messages() {
170 let details = vec![json!({
171 "type": "reasoning.text",
172 "text": "private reasoning"
173 })];
174 let mut assistant = ChatMessage::assistant("answer".to_owned(), Vec::new());
175 assistant.reasoning_details = Some(details.clone());
176 assert_eq!(
177 assistant.to_openai_value()["reasoning_details"],
178 json!(details)
179 );
180
181 let old_message: ChatMessage = serde_json::from_value(json!({
182 "role": "assistant",
183 "content": "old session"
184 }))
185 .expect("old message without reasoning details");
186 assert_eq!(old_message.reasoning_details, None);
187
188 let mut user = ChatMessage::user("question".to_owned());
189 user.reasoning_details = Some(vec![json!({"text": "must not be sent"})]);
190 assert!(user.to_openai_value().get("reasoning_details").is_none());
191 }
192}