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::edit_file_tool(),
49                Self::web_search_tool(),
50                Self::web_fetch_tool(),
51                Self::agent_tool(),
52            ],
53        }
54    }
55
56    /// Get a reference to the cached Ollama tool definitions without constructing a registry
57    pub fn ollama_tools_cached() -> &'static [serde_json::Value] {
58        &OLLAMA_TOOLS_CACHE
59    }
60
61    /// Get all tools
62    pub fn tools(&self) -> &[Tool] {
63        &self.tools
64    }
65
66    // Tool Definitions
67
68    fn read_file_tool() -> Tool {
69        Tool {
70            type_: "function".to_string(),
71            function: ToolFunction {
72                name: "read_file".to_string(),
73                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(),
74                parameters: json!({
75                    "type": "object",
76                    "properties": {
77                        "path": {
78                            "type": "string",
79                            "description": "Absolute or relative path to the file to read. Use absolute paths (e.g., /home/user/file.pdf) for files outside the project."
80                        }
81                    },
82                    "required": ["path"]
83                }),
84            },
85        }
86    }
87
88    fn write_file_tool() -> Tool {
89        Tool {
90            type_: "function".to_string(),
91            function: ToolFunction {
92                name: "write_file".to_string(),
93                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(),
94                parameters: json!({
95                    "type": "object",
96                    "properties": {
97                        "path": {
98                            "type": "string",
99                            "description": "Path to the file to write, relative to the project root or absolute (must be within project)"
100                        },
101                        "content": {
102                            "type": "string",
103                            "description": "The complete file content to write"
104                        }
105                    },
106                    "required": ["path", "content"]
107                }),
108            },
109        }
110    }
111
112    fn delete_file_tool() -> Tool {
113        Tool {
114            type_: "function".to_string(),
115            function: ToolFunction {
116                name: "delete_file".to_string(),
117                description: "Delete a file from the project directory. Creates a timestamped backup before deletion for recovery.".to_string(),
118                parameters: json!({
119                    "type": "object",
120                    "properties": {
121                        "path": {
122                            "type": "string",
123                            "description": "Path to the file to delete"
124                        }
125                    },
126                    "required": ["path"]
127                }),
128            },
129        }
130    }
131
132    fn create_directory_tool() -> Tool {
133        Tool {
134            type_: "function".to_string(),
135            function: ToolFunction {
136                name: "create_directory".to_string(),
137                description:
138                    "Create a new directory in the project. Creates parent directories if needed."
139                        .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 edit_file_tool() -> Tool {
183        Tool {
184            type_: "function".to_string(),
185            function: ToolFunction {
186                name: "edit_file".to_string(),
187                description: "Make targeted edits to a file by replacing specific text. \
188                    The old_string must match exactly and uniquely in the file. \
189                    Prefer this over write_file for modifying existing files."
190                    .to_string(),
191                parameters: json!({
192                    "type": "object",
193                    "properties": {
194                        "path": {
195                            "type": "string",
196                            "description": "Path to the file to edit"
197                        },
198                        "old_string": {
199                            "type": "string",
200                            "description": "The exact text to find and replace (must be unique in the file)"
201                        },
202                        "new_string": {
203                            "type": "string",
204                            "description": "The new text to replace old_string with"
205                        }
206                    },
207                    "required": ["path", "old_string", "new_string"]
208                }),
209            },
210        }
211    }
212
213    fn web_search_tool() -> Tool {
214        Tool {
215            type_: "function".to_string(),
216            function: ToolFunction {
217                name: "web_search".to_string(),
218                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(),
219                parameters: json!({
220                    "type": "object",
221                    "properties": {
222                        "query": {
223                            "type": "string",
224                            "description": "Search query. Be specific and include version numbers when relevant (e.g., 'Rust async tokio 1.40 new features')"
225                        },
226                        "max_results": {
227                            "type": "integer",
228                            "description": "Number of results to fetch (1-10). Use 3 for simple facts, 5-7 for research, 10 for comprehensive analysis.",
229                            "minimum": 1,
230                            "maximum": 10
231                        }
232                    },
233                    "required": ["query", "max_results"]
234                }),
235            },
236        }
237    }
238
239    fn web_fetch_tool() -> Tool {
240        Tool {
241            type_: "function".to_string(),
242            function: ToolFunction {
243                name: "web_fetch".to_string(),
244                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(),
245                parameters: json!({
246                    "type": "object",
247                    "properties": {
248                        "url": {
249                            "type": "string",
250                            "description": "The URL to fetch content from (e.g., 'https://docs.rs/tokio/latest')"
251                        }
252                    },
253                    "required": ["url"]
254                }),
255            },
256        }
257    }
258
259    fn agent_tool() -> Tool {
260        Tool {
261            type_: "function".to_string(),
262            function: ToolFunction {
263                name: "agent".to_string(),
264                description: "Spawn an autonomous sub-agent to handle a task independently. \
265                    The agent gets its own conversation context and full tool access. \
266                    Give it a self-contained task via the prompt parameter. \
267                    Multiple agent calls in one response run in parallel.".to_string(),
268                parameters: json!({
269                    "type": "object",
270                    "properties": {
271                        "prompt": {
272                            "type": "string",
273                            "description": "The task for the agent to complete"
274                        },
275                        "description": {
276                            "type": "string",
277                            "description": "Short label for the UI (e.g., 'Read src/models/ files')"
278                        }
279                    },
280                    "required": ["prompt", "description"]
281                }),
282            },
283        }
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290
291    #[test]
292    fn test_tool_registry_creation() {
293        let registry = ToolRegistry::mermaid_tools();
294        assert_eq!(registry.tools().len(), 9, "Should have 9 tools defined");
295    }
296
297    #[test]
298    fn test_tool_serialization() {
299        let ollama_tools = ToolRegistry::ollama_tools_cached();
300
301        assert_eq!(ollama_tools.len(), 9);
302
303        // Verify first tool has correct structure
304        let first_tool = &ollama_tools[0];
305        assert!(first_tool.get("type").is_some());
306        assert!(first_tool.get("function").is_some());
307    }
308
309    #[test]
310    fn test_read_file_tool_schema() {
311        let tool = ToolRegistry::read_file_tool();
312        assert_eq!(tool.function.name, "read_file");
313        assert!(tool.function.description.contains("Read a file"));
314
315        let params = tool.function.parameters.as_object().unwrap();
316        assert!(params.get("properties").is_some());
317        assert!(params.get("required").is_some());
318    }
319}