Skip to main content

orion_core/
tools.rs

1#[cfg(feature = "tools")]
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4
5#[cfg(feature = "tools")]
6use crate::error::CoreResult;
7
8/// A tool call parsed out of an assistant message.
9///
10/// Produced by [`parse_tool_calls`]. The agent assigns each call an id and
11/// dispatches it to the matching registered `Tool`.
12#[derive(Debug, Clone, PartialEq)]
13pub struct ParsedToolCall {
14    /// Tool name the model wants to invoke.
15    pub name: String,
16    /// Arguments object (defaults to `{}` when the model omits it).
17    pub arguments: serde_json::Value,
18}
19
20/// Parse tool calls out of an assistant message.
21///
22/// Templates advertise the convention rendered by `render_tools`: a fenced
23/// ```` ```tool_call ```` block holding `{"name": ..., "arguments": {...}}`
24/// (or a JSON array of such objects). Parsing is deliberately lenient so
25/// smaller local models still trigger tools when they drift from the exact
26/// format:
27///
28/// 1. Every ```` ```tool_call ```` and ```` ```json ```` fenced block is parsed.
29/// 2. If no fenced block yields a call, the *whole trimmed message* is tried as
30///    a single JSON object — but only when it carries both `name` and
31///    `arguments` keys.
32///
33/// Arbitrary mid-prose substrings are never scanned, so ordinary replies that
34/// merely mention JSON don't produce false positives. Entries without a string
35/// `name` are skipped; a missing `arguments` defaults to an empty object.
36///
37/// ```
38/// use orion_core::parse_tool_calls;
39///
40/// let reply = "Sure.\n```tool_call\n\
41///              {\"name\": \"read_file\", \"arguments\": {\"path\": \"Cargo.toml\"}}\n```";
42/// let calls = parse_tool_calls(reply);
43/// assert_eq!(calls.len(), 1);
44/// assert_eq!(calls[0].name, "read_file");
45/// assert_eq!(calls[0].arguments["path"], "Cargo.toml");
46///
47/// // A plain answer that merely mentions JSON yields nothing.
48/// assert!(parse_tool_calls("Here is a name field somewhere.").is_empty());
49/// ```
50pub fn parse_tool_calls(text: &str) -> Vec<ParsedToolCall> {
51    let mut calls = Vec::new();
52
53    for block in fenced_blocks(text) {
54        if matches!(block.tag.as_str(), "tool_call" | "json") {
55            collect_calls(block.body, &mut calls);
56        }
57    }
58
59    // Fallback: a bare JSON object that is the entire message. Guarded on the
60    // presence of both keys so plain JSON answers aren't mistaken for calls.
61    if calls.is_empty() {
62        let trimmed = text.trim();
63        if trimmed.starts_with('{')
64            && trimmed.contains("\"name\"")
65            && trimmed.contains("\"arguments\"")
66        {
67            collect_calls(trimmed, &mut calls);
68        }
69    }
70
71    calls
72}
73
74/// Parse one JSON snippet (object or array of objects) into tool calls,
75/// appending any well-formed entries to `out`.
76fn collect_calls(snippet: &str, out: &mut Vec<ParsedToolCall>) {
77    let Ok(value) = serde_json::from_str::<serde_json::Value>(snippet.trim()) else {
78        return;
79    };
80    match value {
81        serde_json::Value::Array(items) => {
82            for item in items {
83                if let Some(call) = call_from_value(&item) {
84                    out.push(call);
85                }
86            }
87        }
88        other => {
89            if let Some(call) = call_from_value(&other) {
90                out.push(call);
91            }
92        }
93    }
94}
95
96/// Extract a single [`ParsedToolCall`] from a JSON value, if it names a tool.
97fn call_from_value(value: &serde_json::Value) -> Option<ParsedToolCall> {
98    let name = value.get("name")?.as_str()?.trim().to_string();
99    if name.is_empty() {
100        return None;
101    }
102    let arguments = value
103        .get("arguments")
104        .cloned()
105        .unwrap_or_else(|| serde_json::json!({}));
106    Some(ParsedToolCall { name, arguments })
107}
108
109/// A fenced code block: its info-string tag (lower-cased, may be empty) and body.
110struct FencedBlock<'a> {
111    tag: String,
112    body: &'a str,
113}
114
115/// Yield every ```` ``` ````-fenced code block in `text`. Tolerant of leading
116/// indentation on the fence and of a missing closing fence at end of input.
117fn fenced_blocks(text: &str) -> Vec<FencedBlock<'_>> {
118    let mut blocks = Vec::new();
119    let bytes = text.as_bytes();
120    let mut search = 0;
121
122    while let Some(rel) = text[search..].find("```") {
123        let open = search + rel;
124        // Tag runs from after the fence to the end of that line.
125        let after_fence = open + 3;
126        let line_end = text[after_fence..]
127            .find('\n')
128            .map(|i| after_fence + i)
129            .unwrap_or(bytes.len());
130        let tag = text[after_fence..line_end].trim().to_ascii_lowercase();
131        let body_start = (line_end + 1).min(bytes.len());
132
133        // Body runs to the next closing fence, or to EOF if unterminated.
134        let (body_end, next) = match text[body_start..].find("```") {
135            Some(i) => (body_start + i, body_start + i + 3),
136            None => (bytes.len(), bytes.len()),
137        };
138
139        blocks.push(FencedBlock {
140            tag,
141            body: &text[body_start..body_end],
142        });
143        search = next;
144    }
145
146    blocks
147}
148
149/// Schema describing a tool's parameters (JSON Schema subset).
150#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct ToolSchema {
152    /// Tool name (must match what the model emits in a tool call).
153    pub name: String,
154    /// Human-readable description shown to the model.
155    pub description: String,
156    /// JSON Schema describing the tool's accepted arguments.
157    pub parameters: serde_json::Value,
158}
159
160/// Result returned by a tool execution.
161#[cfg(feature = "tools")]
162#[derive(Debug, Clone, Serialize, Deserialize)]
163pub struct ToolOutput {
164    /// Output content fed back to the model as the tool result.
165    pub content: String,
166    /// Optional structured details for the UI (not sent to the model).
167    #[serde(default)]
168    pub details: serde_json::Value,
169}
170
171/// Callback for streaming tool progress.
172#[cfg(feature = "tools")]
173pub type ToolUpdateCallback = Box<dyn FnMut(&str) + Send>;
174
175/// A tool the agent can invoke.
176///
177/// Tools are defined by the host application and registered
178/// with the agent. The agent loop calls `execute` when the
179/// LLM emits a tool call.
180///
181/// ```no_run
182/// use orion_core::{Tool, ToolOutput, CoreError, CoreResult};
183/// use orion_core::tools::ToolUpdateCallback;
184/// use async_trait::async_trait;
185///
186/// struct ReadFileTool;
187///
188/// #[async_trait]
189/// impl Tool for ReadFileTool {
190///     fn name(&self) -> &str { "read_file" }
191///     fn label(&self) -> &str { "Read File" }
192///     fn description(&self) -> &str { "Read a file from disk" }
193///
194///     fn parameters_schema(&self) -> serde_json::Value {
195///         serde_json::json!({
196///             "type": "object",
197///             "properties": { "path": { "type": "string" } },
198///             "required": ["path"]
199///         })
200///     }
201///
202///     async fn execute(
203///         &self,
204///         _tool_call_id: &str,
205///         args: serde_json::Value,
206///         _on_update: Option<ToolUpdateCallback>,
207///     ) -> CoreResult<ToolOutput> {
208///         let path = args["path"].as_str().unwrap_or("");
209///         let content = std::fs::read_to_string(path)
210///             .map_err(|e| CoreError::Tool(e.to_string()))?;
211///         Ok(ToolOutput { content, details: serde_json::json!({ "path": path }) })
212///     }
213/// }
214///
215/// // Register with the agent: `agent.set_tools(vec![Box::new(ReadFileTool)]);`
216/// ```
217#[cfg(feature = "tools")]
218#[async_trait]
219pub trait Tool: Send + Sync {
220    /// Unique tool name (must match what the LLM outputs).
221    fn name(&self) -> &str;
222
223    /// Human-readable label for UI display.
224    fn label(&self) -> &str;
225
226    /// Description shown to the LLM in the system prompt.
227    fn description(&self) -> &str;
228
229    /// JSON Schema for the tool's parameters.
230    fn parameters_schema(&self) -> serde_json::Value;
231
232    /// Execute the tool with the given arguments.
233    async fn execute(
234        &self,
235        tool_call_id: &str,
236        args: serde_json::Value,
237        on_update: Option<ToolUpdateCallback>,
238    ) -> CoreResult<ToolOutput>;
239
240    /// Return the full schema for system prompt injection.
241    fn schema(&self) -> ToolSchema {
242        ToolSchema {
243            name: self.name().to_string(),
244            description: self.description().to_string(),
245            parameters: self.parameters_schema(),
246        }
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use serde_json::json;
254
255    #[test]
256    fn parses_fenced_tool_call() {
257        let text =
258            "Sure.\n```tool_call\n{\"name\": \"add\", \"arguments\": {\"a\": 2, \"b\": 3}}\n```";
259        let calls = parse_tool_calls(text);
260        assert_eq!(calls.len(), 1);
261        assert_eq!(calls[0].name, "add");
262        assert_eq!(calls[0].arguments, json!({"a": 2, "b": 3}));
263    }
264
265    #[test]
266    fn parses_array_of_calls() {
267        let text = "```tool_call\n[{\"name\": \"a\", \"arguments\": {}}, {\"name\": \"b\", \"arguments\": {\"x\": 1}}]\n```";
268        let calls = parse_tool_calls(text);
269        assert_eq!(calls.len(), 2);
270        assert_eq!(calls[0].name, "a");
271        assert_eq!(calls[1].name, "b");
272        assert_eq!(calls[1].arguments, json!({"x": 1}));
273    }
274
275    #[test]
276    fn parses_json_tagged_fence() {
277        let text = "```json\n{\"name\": \"now\", \"arguments\": {}}\n```";
278        let calls = parse_tool_calls(text);
279        assert_eq!(calls.len(), 1);
280        assert_eq!(calls[0].name, "now");
281    }
282
283    #[test]
284    fn parses_bare_object_whole_message() {
285        let text = "  {\"name\": \"now\", \"arguments\": {}}  ";
286        let calls = parse_tool_calls(text);
287        assert_eq!(calls.len(), 1);
288        assert_eq!(calls[0].name, "now");
289    }
290
291    #[test]
292    fn missing_arguments_defaults_to_empty_object() {
293        let text = "```tool_call\n{\"name\": \"now\"}\n```";
294        let calls = parse_tool_calls(text);
295        assert_eq!(calls.len(), 1);
296        assert_eq!(calls[0].arguments, json!({}));
297    }
298
299    #[test]
300    fn collects_multiple_fenced_blocks() {
301        let text = "```tool_call\n{\"name\": \"a\", \"arguments\": {}}\n```\nthen\n```tool_call\n{\"name\": \"b\", \"arguments\": {}}\n```";
302        let calls = parse_tool_calls(text);
303        assert_eq!(calls.len(), 2);
304    }
305
306    #[test]
307    fn skips_malformed_and_nameless() {
308        let text = "```tool_call\nnot json at all\n```";
309        assert!(parse_tool_calls(text).is_empty());
310        let nameless = "```tool_call\n{\"arguments\": {}}\n```";
311        assert!(parse_tool_calls(nameless).is_empty());
312    }
313
314    #[test]
315    fn plain_prose_yields_no_calls() {
316        let text = "Here is some JSON you might use: it has a name field somewhere.";
317        assert!(parse_tool_calls(text).is_empty());
318        // A normal JSON answer without `arguments` must not be treated as a call.
319        let answer = "{\"name\": \"Ada\", \"age\": 36}";
320        assert!(parse_tool_calls(answer).is_empty());
321    }
322
323    #[test]
324    fn handles_unterminated_fence() {
325        let text = "```tool_call\n{\"name\": \"now\", \"arguments\": {}}";
326        let calls = parse_tool_calls(text);
327        assert_eq!(calls.len(), 1);
328        assert_eq!(calls[0].name, "now");
329    }
330}