Skip to main content

mermaid_cli/models/
tools.rs

1/// Ollama Tools API support for native function calling
2///
3/// This module defines Mermaid's available tools in Ollama's JSON Schema format,
4/// replacing the legacy text-based action block system.
5
6use serde::{Deserialize, Serialize};
7use serde_json::json;
8
9/// A tool available to the model (Ollama format)
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct Tool {
12    #[serde(rename = "type")]
13    pub type_: String,
14    pub function: ToolFunction,
15}
16
17/// Function definition for a tool
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct ToolFunction {
20    pub name: String,
21    pub description: String,
22    pub parameters: serde_json::Value,
23}
24
25/// Registry of all available Mermaid tools
26pub struct ToolRegistry {
27    tools: Vec<Tool>,
28}
29
30impl ToolRegistry {
31    /// Create a new registry with all Mermaid tools
32    pub fn mermaid_tools() -> Self {
33        Self {
34            tools: vec![
35                Self::read_file_tool(),
36                Self::write_file_tool(),
37                Self::delete_file_tool(),
38                Self::create_directory_tool(),
39                Self::execute_command_tool(),
40                Self::git_diff_tool(),
41                Self::git_status_tool(),
42                Self::git_commit_tool(),
43                Self::web_search_tool(),
44            ],
45        }
46    }
47
48    /// Get all tools in Ollama JSON format
49    pub fn to_ollama_format(&self) -> Vec<serde_json::Value> {
50        self.tools.iter().map(|t| json!(t)).collect()
51    }
52
53    /// Get all tools
54    pub fn tools(&self) -> &[Tool] {
55        &self.tools
56    }
57
58    // Tool Definitions
59
60    fn read_file_tool() -> Tool {
61        Tool {
62            type_: "function".to_string(),
63            function: ToolFunction {
64                name: "read_file".to_string(),
65                description: "Read a file from the filesystem. Can read files anywhere on the system the user has access to, including outside the current project directory. Supports text files, PDFs (sent to vision models), and images.".to_string(),
66                parameters: json!({
67                    "type": "object",
68                    "properties": {
69                        "path": {
70                            "type": "string",
71                            "description": "Absolute or relative path to the file to read. Use absolute paths (e.g., /home/user/file.pdf) for files outside the project."
72                        }
73                    },
74                    "required": ["path"]
75                }),
76            },
77        }
78    }
79
80    fn write_file_tool() -> Tool {
81        Tool {
82            type_: "function".to_string(),
83            function: ToolFunction {
84                name: "write_file".to_string(),
85                description: "Write or create a file in the current project directory. Creates parent directories if they don't exist. Creates a timestamped backup if the file already exists.".to_string(),
86                parameters: json!({
87                    "type": "object",
88                    "properties": {
89                        "path": {
90                            "type": "string",
91                            "description": "Path to the file to write, relative to the project root or absolute (must be within project)"
92                        },
93                        "content": {
94                            "type": "string",
95                            "description": "The complete file content to write"
96                        }
97                    },
98                    "required": ["path", "content"]
99                }),
100            },
101        }
102    }
103
104    fn delete_file_tool() -> Tool {
105        Tool {
106            type_: "function".to_string(),
107            function: ToolFunction {
108                name: "delete_file".to_string(),
109                description: "Delete a file from the project directory. Creates a timestamped backup before deletion for recovery.".to_string(),
110                parameters: json!({
111                    "type": "object",
112                    "properties": {
113                        "path": {
114                            "type": "string",
115                            "description": "Path to the file to delete"
116                        }
117                    },
118                    "required": ["path"]
119                }),
120            },
121        }
122    }
123
124    fn create_directory_tool() -> Tool {
125        Tool {
126            type_: "function".to_string(),
127            function: ToolFunction {
128                name: "create_directory".to_string(),
129                description: "Create a new directory in the project. Creates parent directories if needed.".to_string(),
130                parameters: json!({
131                    "type": "object",
132                    "properties": {
133                        "path": {
134                            "type": "string",
135                            "description": "Path to the directory to create"
136                        }
137                    },
138                    "required": ["path"]
139                }),
140            },
141        }
142    }
143
144    fn execute_command_tool() -> Tool {
145        Tool {
146            type_: "function".to_string(),
147            function: ToolFunction {
148                name: "execute_command".to_string(),
149                description: "Execute a shell command. Use for running tests, builds, git operations, or any terminal command.".to_string(),
150                parameters: json!({
151                    "type": "object",
152                    "properties": {
153                        "command": {
154                            "type": "string",
155                            "description": "The shell command to execute (e.g., 'cargo test', 'npm install')"
156                        },
157                        "working_dir": {
158                            "type": "string",
159                            "description": "Optional working directory to run the command in. Defaults to project root."
160                        }
161                    },
162                    "required": ["command"]
163                }),
164            },
165        }
166    }
167
168    fn git_diff_tool() -> Tool {
169        Tool {
170            type_: "function".to_string(),
171            function: ToolFunction {
172                name: "git_diff".to_string(),
173                description: "Show git diff for staged and unstaged changes. Can show diff for specific files or entire repository.".to_string(),
174                parameters: json!({
175                    "type": "object",
176                    "properties": {
177                        "path": {
178                            "type": "string",
179                            "description": "Optional specific file path to show diff for. If omitted, shows diff for entire repository."
180                        }
181                    },
182                    "required": []
183                }),
184            },
185        }
186    }
187
188    fn git_status_tool() -> Tool {
189        Tool {
190            type_: "function".to_string(),
191            function: ToolFunction {
192                name: "git_status".to_string(),
193                description: "Show the current git repository status including staged, unstaged, and untracked files.".to_string(),
194                parameters: json!({
195                    "type": "object",
196                    "properties": {},
197                    "required": []
198                }),
199            },
200        }
201    }
202
203    fn git_commit_tool() -> Tool {
204        Tool {
205            type_: "function".to_string(),
206            function: ToolFunction {
207                name: "git_commit".to_string(),
208                description: "Create a git commit with specified message and files.".to_string(),
209                parameters: json!({
210                    "type": "object",
211                    "properties": {
212                        "message": {
213                            "type": "string",
214                            "description": "Commit message"
215                        },
216                        "files": {
217                            "type": "array",
218                            "items": {
219                                "type": "string"
220                            },
221                            "description": "List of file paths to include in the commit"
222                        }
223                    },
224                    "required": ["message", "files"]
225                }),
226            },
227        }
228    }
229
230    fn web_search_tool() -> Tool {
231        Tool {
232            type_: "function".to_string(),
233            function: ToolFunction {
234                name: "web_search".to_string(),
235                description: "Search the web using local Searxng instance. Returns full page content in markdown format for deep analysis. Use for current information, library documentation, version-specific questions, or any time-sensitive data.".to_string(),
236                parameters: json!({
237                    "type": "object",
238                    "properties": {
239                        "query": {
240                            "type": "string",
241                            "description": "Search query. Be specific and include version numbers when relevant (e.g., 'Rust async tokio 1.40 new features')"
242                        },
243                        "result_count": {
244                            "type": "integer",
245                            "description": "Number of results to fetch (1-10). Use 3 for simple facts, 5-7 for research, 10 for comprehensive analysis.",
246                            "minimum": 1,
247                            "maximum": 10
248                        }
249                    },
250                    "required": ["query", "result_count"]
251                }),
252            },
253        }
254    }
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260
261    #[test]
262    fn test_tool_registry_creation() {
263        let registry = ToolRegistry::mermaid_tools();
264        assert_eq!(registry.tools().len(), 9, "Should have 9 tools defined");
265    }
266
267    #[test]
268    fn test_tool_serialization() {
269        let registry = ToolRegistry::mermaid_tools();
270        let ollama_tools = registry.to_ollama_format();
271
272        assert_eq!(ollama_tools.len(), 9);
273
274        // Verify first tool has correct structure
275        let first_tool = &ollama_tools[0];
276        assert!(first_tool.get("type").is_some());
277        assert!(first_tool.get("function").is_some());
278    }
279
280    #[test]
281    fn test_read_file_tool_schema() {
282        let tool = ToolRegistry::read_file_tool();
283        assert_eq!(tool.function.name, "read_file");
284        assert!(tool.function.description.contains("Read a file"));
285
286        let params = tool.function.parameters.as_object().unwrap();
287        assert!(params.get("properties").is_some());
288        assert!(params.get("required").is_some());
289    }
290}