1use serde::{Deserialize, Serialize};
7use serde_json::json;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct Tool {
12 #[serde(rename = "type")]
13 pub type_: String,
14 pub function: ToolFunction,
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct ToolFunction {
20 pub name: String,
21 pub description: String,
22 pub parameters: serde_json::Value,
23}
24
25pub struct ToolRegistry {
27 tools: Vec<Tool>,
28}
29
30impl ToolRegistry {
31 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 pub fn to_ollama_format(&self) -> Vec<serde_json::Value> {
50 self.tools.iter().map(|t| json!(t)).collect()
51 }
52
53 pub fn tools(&self) -> &[Tool] {
55 &self.tools
56 }
57
58 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 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}