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;
8use std::sync::LazyLock;
9
10/// A tool available to the model (Ollama format)
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct Tool {
13    #[serde(rename = "type")]
14    pub type_: String,
15    pub function: ToolFunction,
16}
17
18/// Function definition for a tool
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct ToolFunction {
21    pub name: String,
22    pub description: String,
23    pub parameters: serde_json::Value,
24}
25
26/// Registry of all available Mermaid tools
27pub struct ToolRegistry {
28    tools: Vec<Tool>,
29}
30
31/// Cached Ollama JSON format for the static tool definitions.
32/// Built once on first access, reused for every chat() call.
33static OLLAMA_TOOLS_CACHE: LazyLock<Vec<serde_json::Value>> = LazyLock::new(|| {
34    let registry = ToolRegistry::mermaid_tools();
35    registry.tools.iter().map(|t| json!(t)).collect()
36});
37
38impl ToolRegistry {
39    /// Create a new registry with all Mermaid tools
40    pub fn mermaid_tools() -> Self {
41        Self {
42            tools: vec![
43                Self::read_file_tool(),
44                Self::write_file_tool(),
45                Self::delete_file_tool(),
46                Self::create_directory_tool(),
47                Self::execute_command_tool(),
48                Self::git_diff_tool(),
49                Self::git_status_tool(),
50                Self::git_commit_tool(),
51                Self::edit_file_tool(),
52                Self::web_search_tool(),
53                Self::web_fetch_tool(),
54            ],
55        }
56    }
57
58    /// Get a reference to the cached Ollama tool definitions without constructing a registry
59    pub fn ollama_tools_cached() -> &'static [serde_json::Value] {
60        &OLLAMA_TOOLS_CACHE
61    }
62
63    /// Get all tools
64    pub fn tools(&self) -> &[Tool] {
65        &self.tools
66    }
67
68    // Tool Definitions
69
70    fn read_file_tool() -> Tool {
71        Tool {
72            type_: "function".to_string(),
73            function: ToolFunction {
74                name: "read_file".to_string(),
75                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(),
76                parameters: json!({
77                    "type": "object",
78                    "properties": {
79                        "path": {
80                            "type": "string",
81                            "description": "Absolute or relative path to the file to read. Use absolute paths (e.g., /home/user/file.pdf) for files outside the project."
82                        }
83                    },
84                    "required": ["path"]
85                }),
86            },
87        }
88    }
89
90    fn write_file_tool() -> Tool {
91        Tool {
92            type_: "function".to_string(),
93            function: ToolFunction {
94                name: "write_file".to_string(),
95                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(),
96                parameters: json!({
97                    "type": "object",
98                    "properties": {
99                        "path": {
100                            "type": "string",
101                            "description": "Path to the file to write, relative to the project root or absolute (must be within project)"
102                        },
103                        "content": {
104                            "type": "string",
105                            "description": "The complete file content to write"
106                        }
107                    },
108                    "required": ["path", "content"]
109                }),
110            },
111        }
112    }
113
114    fn delete_file_tool() -> Tool {
115        Tool {
116            type_: "function".to_string(),
117            function: ToolFunction {
118                name: "delete_file".to_string(),
119                description: "Delete a file from the project directory. Creates a timestamped backup before deletion for recovery.".to_string(),
120                parameters: json!({
121                    "type": "object",
122                    "properties": {
123                        "path": {
124                            "type": "string",
125                            "description": "Path to the file to delete"
126                        }
127                    },
128                    "required": ["path"]
129                }),
130            },
131        }
132    }
133
134    fn create_directory_tool() -> Tool {
135        Tool {
136            type_: "function".to_string(),
137            function: ToolFunction {
138                name: "create_directory".to_string(),
139                description: "Create a new directory in the project. Creates parent directories if needed.".to_string(),
140                parameters: json!({
141                    "type": "object",
142                    "properties": {
143                        "path": {
144                            "type": "string",
145                            "description": "Path to the directory to create"
146                        }
147                    },
148                    "required": ["path"]
149                }),
150            },
151        }
152    }
153
154    fn execute_command_tool() -> Tool {
155        Tool {
156            type_: "function".to_string(),
157            function: ToolFunction {
158                name: "execute_command".to_string(),
159                description: "Execute a shell command. Use for running tests, builds, git operations, or any terminal command. For long-running processes like servers, set a short timeout (e.g., 5) — the process will keep running after timeout.".to_string(),
160                parameters: json!({
161                    "type": "object",
162                    "properties": {
163                        "command": {
164                            "type": "string",
165                            "description": "The shell command to execute (e.g., 'cargo test', 'npm install')"
166                        },
167                        "working_dir": {
168                            "type": "string",
169                            "description": "Optional working directory to run the command in. Defaults to project root."
170                        },
171                        "timeout": {
172                            "type": "integer",
173                            "description": "Timeout in seconds (default: 30, max: 300). For servers/daemons, use a short timeout like 5 since the process continues running after timeout."
174                        }
175                    },
176                    "required": ["command"]
177                }),
178            },
179        }
180    }
181
182    fn git_diff_tool() -> Tool {
183        Tool {
184            type_: "function".to_string(),
185            function: ToolFunction {
186                name: "git_diff".to_string(),
187                description: "Show git diff for staged and unstaged changes. Can show diff for specific files or entire repository.".to_string(),
188                parameters: json!({
189                    "type": "object",
190                    "properties": {
191                        "path": {
192                            "type": "string",
193                            "description": "Optional specific file path to show diff for. If omitted, shows diff for entire repository."
194                        }
195                    },
196                    "required": []
197                }),
198            },
199        }
200    }
201
202    fn git_status_tool() -> Tool {
203        Tool {
204            type_: "function".to_string(),
205            function: ToolFunction {
206                name: "git_status".to_string(),
207                description: "Show the current git repository status including staged, unstaged, and untracked files.".to_string(),
208                parameters: json!({
209                    "type": "object",
210                    "properties": {},
211                    "required": []
212                }),
213            },
214        }
215    }
216
217    fn git_commit_tool() -> Tool {
218        Tool {
219            type_: "function".to_string(),
220            function: ToolFunction {
221                name: "git_commit".to_string(),
222                description: "Create a git commit with specified message and files.".to_string(),
223                parameters: json!({
224                    "type": "object",
225                    "properties": {
226                        "message": {
227                            "type": "string",
228                            "description": "Commit message"
229                        },
230                        "files": {
231                            "type": "array",
232                            "items": {
233                                "type": "string"
234                            },
235                            "description": "List of file paths to include in the commit"
236                        }
237                    },
238                    "required": ["message", "files"]
239                }),
240            },
241        }
242    }
243
244    fn edit_file_tool() -> Tool {
245        Tool {
246            type_: "function".to_string(),
247            function: ToolFunction {
248                name: "edit_file".to_string(),
249                description: "Make targeted edits to a file by replacing specific text. \
250                    The old_string must match exactly and uniquely in the file. \
251                    Prefer this over write_file for modifying existing files.".to_string(),
252                parameters: json!({
253                    "type": "object",
254                    "properties": {
255                        "path": {
256                            "type": "string",
257                            "description": "Path to the file to edit"
258                        },
259                        "old_string": {
260                            "type": "string",
261                            "description": "The exact text to find and replace (must be unique in the file)"
262                        },
263                        "new_string": {
264                            "type": "string",
265                            "description": "The new text to replace old_string with"
266                        }
267                    },
268                    "required": ["path", "old_string", "new_string"]
269                }),
270            },
271        }
272    }
273
274    fn web_search_tool() -> Tool {
275        Tool {
276            type_: "function".to_string(),
277            function: ToolFunction {
278                name: "web_search".to_string(),
279                description: "Search the web for information. 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(),
280                parameters: json!({
281                    "type": "object",
282                    "properties": {
283                        "query": {
284                            "type": "string",
285                            "description": "Search query. Be specific and include version numbers when relevant (e.g., 'Rust async tokio 1.40 new features')"
286                        },
287                        "max_results": {
288                            "type": "integer",
289                            "description": "Number of results to fetch (1-10). Use 3 for simple facts, 5-7 for research, 10 for comprehensive analysis.",
290                            "minimum": 1,
291                            "maximum": 10
292                        }
293                    },
294                    "required": ["query", "max_results"]
295                }),
296            },
297        }
298    }
299
300    fn web_fetch_tool() -> Tool {
301        Tool {
302            type_: "function".to_string(),
303            function: ToolFunction {
304                name: "web_fetch".to_string(),
305                description: "Fetch content from a URL and return it as clean markdown. Use for reading documentation pages, articles, GitHub READMEs, or any web page the user references.".to_string(),
306                parameters: json!({
307                    "type": "object",
308                    "properties": {
309                        "url": {
310                            "type": "string",
311                            "description": "The URL to fetch content from (e.g., 'https://docs.rs/tokio/latest')"
312                        }
313                    },
314                    "required": ["url"]
315                }),
316            },
317        }
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[test]
326    fn test_tool_registry_creation() {
327        let registry = ToolRegistry::mermaid_tools();
328        assert_eq!(registry.tools().len(), 11, "Should have 11 tools defined");
329    }
330
331    #[test]
332    fn test_tool_serialization() {
333        let ollama_tools = ToolRegistry::ollama_tools_cached();
334
335        assert_eq!(ollama_tools.len(), 11);
336
337        // Verify first tool has correct structure
338        let first_tool = &ollama_tools[0];
339        assert!(first_tool.get("type").is_some());
340        assert!(first_tool.get("function").is_some());
341    }
342
343    #[test]
344    fn test_read_file_tool_schema() {
345        let tool = ToolRegistry::read_file_tool();
346        assert_eq!(tool.function.name, "read_file");
347        assert!(tool.function.description.contains("Read a file"));
348
349        let params = tool.function.parameters.as_object().unwrap();
350        assert!(params.get("properties").is_some());
351        assert!(params.get("required").is_some());
352    }
353}