use crate::json_repair::{parse_tolerant_json, JsonRepairError};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolCall {
pub id: String,
#[serde(rename = "type")]
pub tool_type: String,
pub function: FunctionCall,
}
impl ToolCall {
pub fn builder(id: impl Into<String>) -> ToolCallBuilder {
ToolCallBuilder::new(id)
}
pub fn name(&self) -> &str {
&self.function.name
}
pub fn arguments(&self) -> &str {
&self.function.arguments
}
pub fn parse_arguments<T: DeserializeOwned>(&self) -> Result<T, JsonRepairError> {
parse_tolerant_json(&self.function.arguments)
}
}
#[derive(Debug, Clone)]
pub struct ToolCallBuilder {
id: String,
tool_type: String,
function: FunctionCall,
}
impl ToolCallBuilder {
pub fn new(id: impl Into<String>) -> Self {
Self {
id: id.into(),
tool_type: "function".to_string(),
function: FunctionCall {
name: String::new(),
arguments: String::new(),
},
}
}
pub fn name(mut self, name: impl Into<String>) -> Self {
self.function.name = name.into();
self
}
pub fn arguments(mut self, arguments: impl Into<String>) -> Self {
self.function.arguments = arguments.into();
self
}
pub fn build(self) -> ToolCall {
ToolCall {
id: self.id,
tool_type: self.tool_type,
function: self.function,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct FunctionCall {
pub name: String,
pub arguments: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallResult {
pub tool_call_id: String,
pub role: String,
pub content: String,
}
impl ToolCallResult {
pub fn new(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
tool_call_id: tool_call_id.into(),
role: "tool".to_string(),
content: content.into(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::collections::HashMap;
#[test]
fn test_tool_call() {
let call = ToolCall::builder("call_123")
.name("calculator")
.arguments(json!({"expression": "2 + 3"}).to_string())
.build();
assert_eq!(call.id, "call_123");
assert_eq!(call.name(), "calculator");
let args: HashMap<String, String> = call.parse_arguments().unwrap();
assert_eq!(args.get("expression").unwrap(), "2 + 3");
}
#[test]
fn test_parse_arguments_tolerates_messy_llm_json() {
let call = ToolCall::builder("call_456")
.name("weather")
.arguments(r#"{"city": "beijing", "unit": "celsius",} plus extra text"#)
.build();
let args: HashMap<String, String> = call.parse_arguments().unwrap();
assert_eq!(args.get("city").unwrap(), "beijing");
assert_eq!(args.get("unit").unwrap(), "celsius");
}
#[test]
fn test_tool_call_result() {
let result = ToolCallResult::new("call_123", "5");
assert_eq!(result.tool_call_id, "call_123");
assert_eq!(result.role, "tool");
assert_eq!(result.content, "5");
}
}