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