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