codei_llm/
message.rs

1use serde::{Deserialize, Deserializer, Serialize, Serializer};
2
3#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
4#[serde(rename_all = "lowercase")]
5pub enum Role {
6    System,
7    User,
8    Assistant,
9    Tool,
10}
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Message {
14    pub role: Role,
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub content: Option<String>,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub tool_calls: Option<Vec<ToolCall>>,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub tool_call_id: Option<String>,
21}
22
23impl Message {
24    pub fn system(content: impl Into<String>) -> Self {
25        Self {
26            role: Role::System,
27            content: Some(content.into()),
28            tool_calls: None,
29            tool_call_id: None,
30        }
31    }
32
33    pub fn user(content: impl Into<String>) -> Self {
34        Self {
35            role: Role::User,
36            content: Some(content.into()),
37            tool_calls: None,
38            tool_call_id: None,
39        }
40    }
41
42    pub fn assistant(content: String, tool_calls: Option<Vec<ToolCall>>) -> Self {
43        let text = content;
44        Self {
45            role: Role::Assistant,
46            content: if text.is_empty() { None } else { Some(text) },
47            tool_calls,
48            tool_call_id: None,
49        }
50    }
51
52    pub fn assistant_with_calls(tool_calls: Vec<ToolCall>) -> Self {
53        Self {
54            role: Role::Assistant,
55            content: None,
56            tool_calls: Some(tool_calls),
57            tool_call_id: None,
58        }
59    }
60
61    pub fn tool(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
62        Self {
63            role: Role::Tool,
64            content: Some(content.into()),
65            tool_calls: None,
66            tool_call_id: Some(tool_call_id.into()),
67        }
68    }
69}
70
71#[derive(Debug, Clone, PartialEq, Eq)]
72pub struct ToolCall {
73    pub id: String,
74    pub name: String,
75    pub arguments: String,
76}
77
78/// OpenAI chat completions expect `tool_calls` entries shaped as
79/// `{ "id", "type": "function", "function": { "name", "arguments" } }`.
80#[derive(Serialize)]
81struct OpenAiToolCall<'a> {
82    id: &'a str,
83    #[serde(rename = "type")]
84    kind: &'static str,
85    function: OpenAiToolFunction<'a>,
86}
87
88#[derive(Serialize)]
89struct OpenAiToolFunction<'a> {
90    name: &'a str,
91    arguments: &'a str,
92}
93
94impl Serialize for ToolCall {
95    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
96    where
97        S: Serializer,
98    {
99        OpenAiToolCall {
100            id: &self.id,
101            kind: "function",
102            function: OpenAiToolFunction {
103                name: &self.name,
104                arguments: &self.arguments,
105            },
106        }
107        .serialize(serializer)
108    }
109}
110
111impl<'de> Deserialize<'de> for ToolCall {
112    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
113    where
114        D: Deserializer<'de>,
115    {
116        #[derive(Deserialize)]
117        struct Raw {
118            id: String,
119            #[serde(default)]
120            name: Option<String>,
121            #[serde(default)]
122            arguments: Option<String>,
123            #[serde(default)]
124            function: Option<OpenAiToolFunctionOwned>,
125        }
126
127        #[derive(Deserialize)]
128        struct OpenAiToolFunctionOwned {
129            name: String,
130            arguments: String,
131        }
132
133        let raw = Raw::deserialize(deserializer)?;
134        if let Some(function) = raw.function {
135            return Ok(Self {
136                id: raw.id,
137                name: function.name,
138                arguments: function.arguments,
139            });
140        }
141
142        Ok(Self {
143            id: raw.id,
144            name: raw
145                .name
146                .ok_or_else(|| serde::de::Error::custom("tool call missing name/function"))?,
147            arguments: raw.arguments.unwrap_or_default(),
148        })
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use serde_json::json;
156
157    #[test]
158    fn serializes_tool_call_for_openai() {
159        let call = ToolCall {
160            id: "call_1".into(),
161            name: "list_dir".into(),
162            arguments: "{}".into(),
163        };
164        let value = serde_json::to_value(&call).unwrap();
165        assert_eq!(
166            value,
167            json!({
168                "id": "call_1",
169                "type": "function",
170                "function": {
171                    "name": "list_dir",
172                    "arguments": "{}"
173                }
174            })
175        );
176    }
177
178    #[test]
179    fn deserializes_openai_and_flat_tool_call_shapes() {
180        let openai = json!({
181            "id": "call_1",
182            "type": "function",
183            "function": {
184                "name": "read",
185                "arguments": r#"{"path":"a.rs"}"#
186            }
187        });
188        let flat = json!({
189            "id": "call_2",
190            "name": "grep",
191            "arguments": r#"{"pattern":"foo"}"#
192        });
193
194        let from_openai: ToolCall = serde_json::from_value(openai).unwrap();
195        let from_flat: ToolCall = serde_json::from_value(flat).unwrap();
196
197        assert_eq!(from_openai.name, "read");
198        assert_eq!(from_flat.name, "grep");
199    }
200}
codei_llm/message.rs

codei_llm/
message.rs
