Skip to main content

lucy/
model.rs

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
26/// Estimate the number of context tokens represented by provider messages.
27///
28/// Lucy supports arbitrary OpenAI-compatible providers and does not bundle a
29/// provider-specific tokenizer. Four UTF-8 bytes per token is therefore a
30/// deliberately conservative display estimate; the statusline should expose
31/// context pressure without pretending that every provider uses the same
32/// tokenizer.
33pub(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}