Skip to main content

tiny_loop/types/
tool.rs

1use schemars::{JsonSchema, generate::SchemaSettings};
2use serde::Serialize;
3use serde_json::{Map, Value};
4
5/// Tool definition for LLM
6#[derive(Serialize, Clone, Debug)]
7pub struct ToolDefinition {
8    /// Type of the tool (typically "function")
9    #[serde(rename = "type")]
10    pub tool_type: String,
11    /// Function definition
12    pub function: ToolFunction,
13}
14
15/// Tool function definition
16#[derive(Serialize, Clone, Debug)]
17pub struct ToolFunction {
18    /// Function name
19    pub name: String,
20    /// Function description
21    pub description: String,
22    /// JSON schema for function parameters
23    pub parameters: Parameters,
24}
25
26/// JSON schema parameters with metadata stripped
27#[derive(Serialize, Clone, Debug)]
28pub struct Parameters(Map<String, Value>);
29
30impl Parameters {
31    /// Create Parameters from a Json object (map)
32    pub fn from_object(mut obj: Map<String, Value>) -> Self {
33        // Remove `$schema`, `title`, and `description` fields from JSON schema
34        obj.remove("$schema");
35        obj.remove("title");
36        obj.remove("description");
37
38        Self(obj)
39    }
40
41    /// Create Parameters from a JsonSchema
42    pub fn from_schema(schema: schemars::Schema) -> Self {
43        let obj = schema.to_value().as_object().unwrap().clone();
44        Self::from_object(obj)
45    }
46
47    /// Create Parameters from a type implementing JsonSchema
48    pub fn from_type<T: JsonSchema>() -> Self {
49        let settings = SchemaSettings::default().with(|s| {
50            s.inline_subschemas = true;
51        });
52        let generator = settings.into_generator();
53        let schema = generator.into_root_schema_for::<T>();
54        Self::from_schema(schema)
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61
62    #[test]
63    fn test_tool_definition_serialization() {
64        let td = ToolDefinition {
65            tool_type: "function".into(),
66            function: ToolFunction {
67                name: "test".into(),
68                description: "desc".into(),
69                parameters: Parameters::from_type::<String>(),
70            },
71        };
72        let json = serde_json::to_string(&td).unwrap();
73        assert!(json.contains(r#""type":"function"#));
74        assert!(json.contains(r#""name":"test"#));
75    }
76}