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
26impl ChatMessage {
27    pub fn system(content: String) -> Self {
28        Self {
29            role: "system".to_owned(),
30            content: Some(content),
31            reasoning_details: None,
32            name: None,
33            tool_call_id: None,
34            tool_calls: Vec::new(),
35        }
36    }
37
38    pub fn user(content: String) -> Self {
39        Self {
40            role: "user".to_owned(),
41            content: Some(content),
42            reasoning_details: None,
43            name: None,
44            tool_call_id: None,
45            tool_calls: Vec::new(),
46        }
47    }
48
49    pub fn assistant(content: String, tool_calls: Vec<ChatToolCall>) -> Self {
50        Self {
51            role: "assistant".to_owned(),
52            content: (!content.is_empty()).then_some(content),
53            reasoning_details: None,
54            name: None,
55            tool_call_id: None,
56            tool_calls,
57        }
58    }
59
60    pub fn tool(tool_call_id: String, name: String, content: String) -> Self {
61        Self {
62            role: "tool".to_owned(),
63            content: Some(content),
64            reasoning_details: None,
65            name: Some(name),
66            tool_call_id: Some(tool_call_id),
67            tool_calls: Vec::new(),
68        }
69    }
70
71    pub fn to_openai_value(&self) -> Value {
72        let mut message = json!({
73            "role": self.role,
74            "content": self.content,
75        });
76        if self.role == "assistant" {
77            if let Some(reasoning_details) = &self.reasoning_details {
78                message["reasoning_details"] = Value::Array(reasoning_details.clone());
79            }
80        }
81        if let Some(name) = &self.name {
82            message["name"] = Value::String(name.clone());
83        }
84        if let Some(tool_call_id) = &self.tool_call_id {
85            message["tool_call_id"] = Value::String(tool_call_id.clone());
86        }
87        if !self.tool_calls.is_empty() {
88            message["tool_calls"] = Value::Array(
89                self.tool_calls
90                    .iter()
91                    .map(|tool_call| {
92                        json!({
93                            "id": tool_call.id,
94                            "type": "function",
95                            "function": {
96                                "name": tool_call.name,
97                                "arguments": tool_call.arguments,
98                            }
99                        })
100                    })
101                    .collect(),
102            );
103        }
104        message
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn tool_assistant_messages_have_openai_compatible_shape() {
114        let assistant = ChatMessage::assistant(
115            String::new(),
116            vec![ChatToolCall {
117                id: "call-1".to_owned(),
118                name: "cmd".to_owned(),
119                arguments: r#"{"command":"pwd"}"#.to_owned(),
120            }],
121        );
122        assert_eq!(assistant.to_openai_value()["content"], Value::Null);
123        assert_eq!(
124            assistant.to_openai_value()["tool_calls"][0]["type"],
125            "function"
126        );
127
128        let tool = ChatMessage::tool(
129            "call-1".to_owned(),
130            "cmd".to_owned(),
131            "{\"exit_code\":0}".to_owned(),
132        );
133        assert_eq!(tool.to_openai_value()["tool_call_id"], "call-1");
134        assert_eq!(tool.to_openai_value()["name"], "cmd");
135    }
136
137    #[test]
138    fn reasoning_details_are_optional_and_only_sent_for_assistant_messages() {
139        let details = vec![json!({
140            "type": "reasoning.text",
141            "text": "private reasoning"
142        })];
143        let mut assistant = ChatMessage::assistant("answer".to_owned(), Vec::new());
144        assistant.reasoning_details = Some(details.clone());
145        assert_eq!(
146            assistant.to_openai_value()["reasoning_details"],
147            json!(details)
148        );
149
150        let old_message: ChatMessage = serde_json::from_value(json!({
151            "role": "assistant",
152            "content": "old session"
153        }))
154        .expect("old message without reasoning details");
155        assert_eq!(old_message.reasoning_details, None);
156
157        let mut user = ChatMessage::user("question".to_owned());
158        user.reasoning_details = Some(vec![json!({"text": "must not be sent"})]);
159        assert!(user.to_openai_value().get("reasoning_details").is_none());
160    }
161}