1use crate::json_repair::{parse_tolerant_json, JsonRepairError};
9use serde::{de::DeserializeOwned, Deserialize, Serialize};
10
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
15pub struct ToolCall {
16 pub id: String,
18
19 #[serde(rename = "type")]
21 pub tool_type: String,
22
23 pub function: FunctionCall,
25}
26
27impl ToolCall {
28 pub fn builder(id: impl Into<String>) -> ToolCallBuilder {
30 ToolCallBuilder::new(id)
31 }
32
33 pub fn name(&self) -> &str {
35 &self.function.name
36 }
37
38 pub fn arguments(&self) -> &str {
40 &self.function.arguments
41 }
42
43 pub fn parse_arguments<T: DeserializeOwned>(&self) -> Result<T, JsonRepairError> {
49 parse_tolerant_json(&self.function.arguments)
50 }
51}
52
53#[derive(Debug, Clone)]
70pub struct ToolCallBuilder {
71 id: String,
72 tool_type: String,
73 function: FunctionCall,
74}
75
76impl ToolCallBuilder {
77 pub fn new(id: impl Into<String>) -> Self {
79 Self {
80 id: id.into(),
81 tool_type: "function".to_string(),
82 function: FunctionCall {
83 name: String::new(),
84 arguments: String::new(),
85 },
86 }
87 }
88
89 pub fn name(mut self, name: impl Into<String>) -> Self {
91 self.function.name = name.into();
92 self
93 }
94
95 pub fn arguments(mut self, arguments: impl Into<String>) -> Self {
97 self.function.arguments = arguments.into();
98 self
99 }
100
101 pub fn build(self) -> ToolCall {
103 ToolCall {
104 id: self.id,
105 tool_type: self.tool_type,
106 function: self.function,
107 }
108 }
109}
110
111#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
113pub struct FunctionCall {
114 pub name: String,
116
117 pub arguments: String,
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct ToolCallResult {
124 pub tool_call_id: String,
126
127 pub role: String,
129
130 pub content: String,
132}
133
134impl ToolCallResult {
135 pub fn new(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
137 Self {
138 tool_call_id: tool_call_id.into(),
139 role: "tool".to_string(),
140 content: content.into(),
141 }
142 }
143}
144
145#[cfg(test)]
146mod tests {
147 use super::*;
148 use serde_json::json;
149 use std::collections::HashMap;
150
151 #[test]
152 fn test_tool_call() {
153 let call = ToolCall::builder("call_123")
154 .name("calculator")
155 .arguments(json!({"expression": "2 + 3"}).to_string())
156 .build();
157
158 assert_eq!(call.id, "call_123");
159 assert_eq!(call.name(), "calculator");
160
161 let args: HashMap<String, String> = call.parse_arguments().unwrap();
162 assert_eq!(args.get("expression").unwrap(), "2 + 3");
163 }
164
165 #[test]
166 fn test_parse_arguments_tolerates_messy_llm_json() {
167 let call = ToolCall::builder("call_456")
169 .name("weather")
170 .arguments(r#"{"city": "beijing", "unit": "celsius",} plus extra text"#)
171 .build();
172
173 let args: HashMap<String, String> = call.parse_arguments().unwrap();
174 assert_eq!(args.get("city").unwrap(), "beijing");
175 assert_eq!(args.get("unit").unwrap(), "celsius");
176 }
177
178 #[test]
179 fn test_tool_call_result() {
180 let result = ToolCallResult::new("call_123", "5");
181
182 assert_eq!(result.tool_call_id, "call_123");
183 assert_eq!(result.role, "tool");
184 assert_eq!(result.content, "5");
185 }
186}