Skip to main content

codetether_agent/rlm/
tools.rs

1//! RLM REPL operations expressed as tool definitions.
2//!
3//! When FunctionGemma is active the RLM loop sends these definitions alongside
4//! the analysis prompt.  The primary LLM (or FunctionGemma after reformatting)
5//! returns structured `ContentPart::ToolCall` entries that are dispatched here
6//! instead of being regex-parsed from code blocks.
7//!
8//! Each tool mirrors a command in the existing DSL REPL (`head`, `tail`,
9//! `grep`, `count`, `llm_query`, `FINAL`).
10
11use crate::provider::ToolDefinition;
12
13/// All RLM REPL operations as tool definitions.
14pub fn rlm_tool_definitions() -> Vec<ToolDefinition> {
15    vec![
16        ToolDefinition {
17            name: "rlm_head".to_string(),
18            description: "Return the first N lines of the loaded context.".to_string(),
19            parameters: serde_json::json!({
20                "type": "object",
21                "properties": {
22                    "n": {
23                        "type": "integer",
24                        "description": "Number of lines from the start (default: 10)"
25                    }
26                },
27                "required": []
28            }),
29        },
30        ToolDefinition {
31            name: "rlm_tail".to_string(),
32            description: "Return the last N lines of the loaded context.".to_string(),
33            parameters: serde_json::json!({
34                "type": "object",
35                "properties": {
36                    "n": {
37                        "type": "integer",
38                        "description": "Number of lines from the end (default: 10)"
39                    }
40                },
41                "required": []
42            }),
43        },
44        ToolDefinition {
45            name: "rlm_grep".to_string(),
46            description: "Search the loaded context for lines matching a regex pattern. Returns matching lines with line numbers.".to_string(),
47            parameters: serde_json::json!({
48                "type": "object",
49                "properties": {
50                    "pattern": {
51                        "type": "string",
52                        "description": "Regex pattern to search for"
53                    }
54                },
55                "required": ["pattern"]
56            }),
57        },
58        ToolDefinition {
59            name: "rlm_count".to_string(),
60            description: "Count occurrences of a regex pattern in the loaded context.".to_string(),
61            parameters: serde_json::json!({
62                "type": "object",
63                "properties": {
64                    "pattern": {
65                        "type": "string",
66                        "description": "Regex pattern to count"
67                    }
68                },
69                "required": ["pattern"]
70            }),
71        },
72        ToolDefinition {
73            name: "rlm_slice".to_string(),
74            description: "Return a slice of the context by line range.".to_string(),
75            parameters: serde_json::json!({
76                "type": "object",
77                "properties": {
78                    "start": {
79                        "type": "integer",
80                        "description": "Start line number (0-indexed)"
81                    },
82                    "end": {
83                        "type": "integer",
84                        "description": "End line number (exclusive)"
85                    }
86                },
87                "required": ["start", "end"]
88            }),
89        },
90        ToolDefinition {
91            name: "rlm_llm_query".to_string(),
92            description: "Ask a focused sub-question about a portion of the context. Use this for semantic understanding of specific sections.".to_string(),
93            parameters: serde_json::json!({
94                "type": "object",
95                "properties": {
96                    "query": {
97                        "type": "string",
98                        "description": "The question to answer about the context"
99                    },
100                    "context_slice": {
101                        "type": "string",
102                        "description": "Optional: specific text slice to analyze (if omitted, uses full context)"
103                    }
104                },
105                "required": ["query"]
106            }),
107        },
108        ToolDefinition {
109            name: "rlm_final".to_string(),
110            description: "Return the final answer to the analysis query. Call this when you have gathered enough information to answer.".to_string(),
111            parameters: serde_json::json!({
112                "type": "object",
113                "properties": {
114                    "answer": {
115                        "type": "string",
116                        "description": "The complete, detailed answer to the original query"
117                    }
118                },
119                "required": ["answer"]
120            }),
121        },
122    ]
123}
124
125/// Result of dispatching an RLM tool call.
126pub enum RlmToolResult {
127    /// Normal output to feed back to the LLM.
128    Output(String),
129    /// The final answer — terminates the RLM loop.
130    Final(String),
131}
132
133/// Dispatch a structured tool call against the REPL.
134///
135/// Returns `None` if the tool name is not an `rlm_*` tool (pass-through for
136/// any other tool calls the model may have produced).
137pub fn dispatch_tool_call(
138    name: &str,
139    arguments: &str,
140    repl: &mut super::repl::RlmRepl,
141) -> Option<RlmToolResult> {
142    let args: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
143
144    match name {
145        "rlm_head" => {
146            let n = args.get("n").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
147            let output = repl.head(n).join("\n");
148            Some(RlmToolResult::Output(output))
149        }
150        "rlm_tail" => {
151            let n = args.get("n").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
152            let output = repl.tail(n).join("\n");
153            Some(RlmToolResult::Output(output))
154        }
155        "rlm_grep" => {
156            let pattern = args.get("pattern").and_then(|v| v.as_str()).unwrap_or("");
157            let matches = repl.grep(pattern);
158            let output = matches
159                .iter()
160                .map(|(i, line)| format!("{}:{}", i, line))
161                .collect::<Vec<_>>()
162                .join("\n");
163            if output.is_empty() {
164                Some(RlmToolResult::Output("(no matches)".to_string()))
165            } else {
166                Some(RlmToolResult::Output(output))
167            }
168        }
169        "rlm_count" => {
170            let pattern = args.get("pattern").and_then(|v| v.as_str()).unwrap_or("");
171            let count = repl.count(pattern);
172            Some(RlmToolResult::Output(count.to_string()))
173        }
174        "rlm_slice" => {
175            let start = args.get("start").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
176            let end = args.get("end").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
177            let output = repl.slice(start, end).to_string();
178            Some(RlmToolResult::Output(output))
179        }
180        "rlm_llm_query" => {
181            // The llm_query tool requires async provider calls — return a
182            // sentinel so the caller knows to handle it specially.
183            let query = args
184                .get("query")
185                .and_then(|v| v.as_str())
186                .unwrap_or("")
187                .to_string();
188            let context_slice = args
189                .get("context_slice")
190                .and_then(|v| v.as_str())
191                .map(|s| s.to_string());
192            // Encode the query + optional slice as JSON so the caller can
193            // destructure it.
194            let payload = serde_json::json!({
195                "__rlm_llm_query": true,
196                "query": query,
197                "context_slice": context_slice,
198            });
199            Some(RlmToolResult::Output(payload.to_string()))
200        }
201        "rlm_final" => {
202            let answer = args
203                .get("answer")
204                .and_then(|v| v.as_str())
205                .unwrap_or("")
206                .to_string();
207            Some(RlmToolResult::Final(answer))
208        }
209        _ => None, // Not an RLM tool
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use crate::rlm::repl::{ReplRuntime, RlmRepl};
217
218    #[test]
219    fn tool_definitions_are_complete() {
220        let defs = rlm_tool_definitions();
221        assert_eq!(defs.len(), 7);
222        let names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
223        assert!(names.contains(&"rlm_head"));
224        assert!(names.contains(&"rlm_tail"));
225        assert!(names.contains(&"rlm_grep"));
226        assert!(names.contains(&"rlm_count"));
227        assert!(names.contains(&"rlm_slice"));
228        assert!(names.contains(&"rlm_llm_query"));
229        assert!(names.contains(&"rlm_final"));
230    }
231
232    #[test]
233    fn dispatch_head() {
234        let ctx = "line 1\nline 2\nline 3\nline 4\nline 5".to_string();
235        let mut repl = RlmRepl::new(ctx, ReplRuntime::Rust);
236        let result = dispatch_tool_call("rlm_head", r#"{"n": 2}"#, &mut repl);
237        match result {
238            Some(RlmToolResult::Output(s)) => assert_eq!(s, "line 1\nline 2"),
239            _ => panic!("expected Output"),
240        }
241    }
242
243    #[test]
244    fn dispatch_tail() {
245        let ctx = "line 1\nline 2\nline 3\nline 4\nline 5".to_string();
246        let mut repl = RlmRepl::new(ctx, ReplRuntime::Rust);
247        let result = dispatch_tool_call("rlm_tail", r#"{"n": 2}"#, &mut repl);
248        match result {
249            Some(RlmToolResult::Output(s)) => assert_eq!(s, "line 4\nline 5"),
250            _ => panic!("expected Output"),
251        }
252    }
253
254    #[test]
255    fn dispatch_grep() {
256        let ctx = "error: fail\ninfo: ok\nerror: boom".to_string();
257        let mut repl = RlmRepl::new(ctx, ReplRuntime::Rust);
258        let result = dispatch_tool_call("rlm_grep", r#"{"pattern": "error"}"#, &mut repl);
259        match result {
260            Some(RlmToolResult::Output(s)) => {
261                assert!(s.contains("error: fail"));
262                assert!(s.contains("error: boom"));
263            }
264            _ => panic!("expected Output"),
265        }
266    }
267
268    #[test]
269    fn dispatch_final() {
270        let ctx = "whatever".to_string();
271        let mut repl = RlmRepl::new(ctx, ReplRuntime::Rust);
272        let result =
273            dispatch_tool_call("rlm_final", r#"{"answer": "The answer is 42"}"#, &mut repl);
274        match result {
275            Some(RlmToolResult::Final(s)) => assert_eq!(s, "The answer is 42"),
276            _ => panic!("expected Final"),
277        }
278    }
279
280    #[test]
281    fn dispatch_unknown_returns_none() {
282        let ctx = "data".to_string();
283        let mut repl = RlmRepl::new(ctx, ReplRuntime::Rust);
284        assert!(dispatch_tool_call("unknown_tool", "{}", &mut repl).is_none());
285    }
286}