pe_tools/
tool.rs

1//! Tool trait — the interface every callable tool implements.
2//!
3//! The LLM sees tools via [`ToolSchema`] (name, description, JSON schema).
4//! When the LLM emits a `ToolCall`, the engine looks up the tool by name
5//! and calls [`Tool::execute`] with the args.
6//!
7//! Two implementations ship with the library:
8//! - [`FunctionTool`] — wraps any async function as a tool
9//! - MCP adapter (future plan) — bridges external MCP servers
10
11use pe_core::error::PeError;
12use pe_core::llm::ToolSchema;
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use std::collections::HashMap;
16use std::future::Future;
17use std::pin::Pin;
18use std::sync::Arc;
19
20/// A callable tool the LLM can invoke via tool_calls.
21///
22/// Implementations: [`FunctionTool`] (simple fns), MCP adapter (future),
23/// graph-as-tool (Plan 014), handoff tools (Plan 009).
24///
25/// # Example
26///
27/// ```ignore
28/// struct WeatherTool;
29///
30/// impl Tool for WeatherTool {
31///     fn name(&self) -> &str { "get_weather" }
32///     fn description(&self) -> &str { "Get current weather for a city" }
33///     fn schema(&self) -> ToolSchema {
34///         ToolSchema {
35///             name: "get_weather".into(),
36///             description: "Get current weather for a city".into(),
37///             parameters: serde_json::json!({
38///                 "type": "object",
39///                 "properties": { "city": { "type": "string" } },
40///                 "required": ["city"]
41///             }),
42///             strict: false,
43///         }
44///     }
45///     fn execute(&self, input: Value) -> ToolFuture {
46///         Box::pin(async move {
47///             Ok(serde_json::json!({ "temp": 72, "unit": "F" }))
48///         })
49///     }
50/// }
51/// ```
52pub trait Tool: Send + Sync {
53    /// Tool name — must be unique within a [`super::ToolRegistry`].
54    /// Sent to the LLM in the tools list.
55    fn name(&self) -> &str;
56
57    /// Human-readable description of what this tool does.
58    /// The LLM reads this to decide when to call the tool.
59    fn description(&self) -> &str;
60
61    /// JSON Schema for this tool's input parameters.
62    /// [`InjectedState`](super::InjectedState) and
63    /// [`InjectedStore`](super::InjectedStore) parameters MUST be excluded.
64    fn schema(&self) -> ToolSchema;
65
66    /// Execute the tool with the given input. Called by [`super::ToolNode`].
67    fn execute(&self, input: Value) -> ToolFuture;
68
69    /// Execute with structured result metadata.
70    ///
71    /// Default implementation wraps `execute()` in `ToolResult::ok()`.
72    /// Override to add metadata (result count, source, confidence, etc.).
73    fn execute_structured(&self, input: Value) -> ToolResultFuture {
74        let fut = self.execute(input);
75        Box::pin(async move {
76            let output = fut.await?;
77            Ok(ToolResult::ok(output))
78        })
79    }
80}
81
82/// Type-erased async future returned by [`Tool::execute`].
83pub type ToolFuture = Pin<Box<dyn Future<Output = Result<Value, PeError>> + Send>>;
84
85/// Type-erased async future for structured tool results.
86pub type ToolResultFuture = Pin<Box<dyn Future<Output = Result<ToolResult, PeError>> + Send>>;
87
88/// Structured result from tool execution.
89///
90/// Wraps the tool's output value with optional metadata (result count, source,
91/// confidence, execution details, etc.). Backward-compatible: `Value` auto-converts
92/// to `ToolResult::ok(value)`.
93///
94/// # Example
95///
96/// ```
97/// use pe_tools::tool::ToolResult;
98/// use serde_json::json;
99///
100/// let result = ToolResult::ok(json!({"answer": 42}))
101///     .with_metadata("source", json!("database"))
102///     .with_metadata("result_count", json!(1));
103/// assert!(result.success);
104/// assert_eq!(result.metadata["source"], "database");
105/// ```
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct ToolResult {
108    /// The tool's output value.
109    pub output: Value,
110    /// Whether the tool succeeded.
111    pub success: bool,
112    /// Tool-specific metadata (result count, source, confidence, etc.).
113    #[serde(default)]
114    pub metadata: HashMap<String, Value>,
115}
116
117impl ToolResult {
118    /// Create a successful result.
119    pub fn ok(output: Value) -> Self {
120        Self {
121            output,
122            success: true,
123            metadata: HashMap::new(),
124        }
125    }
126
127    /// Create an error result with a message.
128    pub fn error(msg: impl Into<String>) -> Self {
129        Self {
130            output: Value::String(msg.into()),
131            success: false,
132            metadata: HashMap::new(),
133        }
134    }
135
136    /// Add metadata to this result.
137    #[must_use]
138    pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
139        self.metadata.insert(key.into(), value);
140        self
141    }
142}
143
144/// Backward compatibility: raw Value auto-converts to successful ToolResult.
145impl From<Value> for ToolResult {
146    fn from(v: Value) -> Self {
147        Self::ok(v)
148    }
149}
150
151/// Type alias for the async function signature used by [`FunctionTool`].
152pub type ToolFunc = Arc<
153    dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, PeError>> + Send>> + Send + Sync,
154>;
155
156/// Convenience wrapper: turns any async function into a [`Tool`].
157///
158/// # Example
159///
160/// ```ignore
161/// let tool = FunctionTool::new(
162///     "add",
163///     "Add two numbers",
164///     serde_json::json!({
165///         "type": "object",
166///         "properties": {
167///             "a": { "type": "number" },
168///             "b": { "type": "number" }
169///         },
170///         "required": ["a", "b"]
171///     }),
172///     |input| Box::pin(async move {
173///         let a = input["a"].as_f64().unwrap_or(0.0);
174///         let b = input["b"].as_f64().unwrap_or(0.0);
175///         Ok(serde_json::json!(a + b))
176///     }),
177/// );
178/// ```
179pub struct FunctionTool {
180    name: String,
181    description: String,
182    schema: ToolSchema,
183    func: ToolFunc,
184}
185
186impl FunctionTool {
187    /// Create a new function tool.
188    ///
189    /// `parameters` is the JSON Schema for the tool's input — passed to the LLM.
190    pub fn new(
191        name: impl Into<String>,
192        description: impl Into<String>,
193        parameters: Value,
194        func: impl Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, PeError>> + Send>>
195        + Send
196        + Sync
197        + 'static,
198    ) -> Self {
199        let name = name.into();
200        let description = description.into();
201        Self {
202            schema: ToolSchema {
203                name: name.clone(),
204                description: description.clone(),
205                parameters,
206                strict: false,
207            },
208            name,
209            description,
210            func: Arc::new(func),
211        }
212    }
213}
214
215impl Tool for FunctionTool {
216    fn name(&self) -> &str {
217        &self.name
218    }
219
220    fn description(&self) -> &str {
221        &self.description
222    }
223
224    fn schema(&self) -> ToolSchema {
225        self.schema.clone()
226    }
227
228    fn execute(&self, input: Value) -> ToolFuture {
229        (self.func)(input)
230    }
231}
232
233/// Type alias for the async function returning a structured [`ToolResult`].
234pub type StructuredToolFunc = Arc<
235    dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<ToolResult, PeError>> + Send>>
236        + Send
237        + Sync,
238>;
239
240/// Like [`FunctionTool`], but returns structured [`ToolResult`] with metadata.
241///
242/// Use this when your tool needs to report result count, source, confidence, etc.
243pub struct StructuredFunctionTool {
244    name: String,
245    description: String,
246    schema: ToolSchema,
247    func: StructuredToolFunc,
248}
249
250impl StructuredFunctionTool {
251    /// Create a new structured function tool.
252    pub fn new(
253        name: impl Into<String>,
254        description: impl Into<String>,
255        parameters: Value,
256        func: impl Fn(Value) -> Pin<Box<dyn Future<Output = Result<ToolResult, PeError>> + Send>>
257        + Send
258        + Sync
259        + 'static,
260    ) -> Self {
261        let name = name.into();
262        let description = description.into();
263        Self {
264            schema: ToolSchema {
265                name: name.clone(),
266                description: description.clone(),
267                parameters,
268                strict: false,
269            },
270            name,
271            description,
272            func: Arc::new(func),
273        }
274    }
275}
276
277impl Tool for StructuredFunctionTool {
278    fn name(&self) -> &str {
279        &self.name
280    }
281
282    fn description(&self) -> &str {
283        &self.description
284    }
285
286    fn schema(&self) -> ToolSchema {
287        self.schema.clone()
288    }
289
290    fn execute(&self, input: Value) -> ToolFuture {
291        let func = self.func.clone();
292        Box::pin(async move {
293            let result = func(input).await?;
294            Ok(result.output)
295        })
296    }
297
298    fn execute_structured(&self, input: Value) -> ToolResultFuture {
299        (self.func)(input)
300    }
301}
302
303impl std::fmt::Debug for StructuredFunctionTool {
304    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
305        f.debug_struct("StructuredFunctionTool")
306            .field("name", &self.name)
307            .field("description", &self.description)
308            .finish()
309    }
310}
311
312impl std::fmt::Debug for FunctionTool {
313    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
314        f.debug_struct("FunctionTool")
315            .field("name", &self.name)
316            .field("description", &self.description)
317            .finish()
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[tokio::test]
326    async fn function_tool_executes_correctly() {
327        let tool = FunctionTool::new(
328            "add",
329            "Add two numbers",
330            serde_json::json!({
331                "type": "object",
332                "properties": {
333                    "a": { "type": "number" },
334                    "b": { "type": "number" }
335                },
336                "required": ["a", "b"]
337            }),
338            |input| {
339                Box::pin(async move {
340                    let a = input["a"].as_f64().unwrap_or(0.0);
341                    let b = input["b"].as_f64().unwrap_or(0.0);
342                    Ok(serde_json::json!(a + b))
343                })
344            },
345        );
346
347        assert_eq!(tool.name(), "add");
348        assert_eq!(tool.description(), "Add two numbers");
349
350        let schema = tool.schema();
351        assert_eq!(schema.name, "add");
352        assert!(!schema.strict);
353
354        let result = tool
355            .execute(serde_json::json!({"a": 3, "b": 4}))
356            .await
357            .unwrap();
358        assert_eq!(result, serde_json::json!(7.0));
359    }
360
361    #[test]
362    fn test_tool_result_ok() {
363        let result = ToolResult::ok(serde_json::json!({"answer": 42}));
364        assert!(result.success);
365        assert_eq!(result.output, serde_json::json!({"answer": 42}));
366        assert!(result.metadata.is_empty());
367    }
368
369    #[test]
370    fn test_tool_result_error() {
371        let result = ToolResult::error("something went wrong");
372        assert!(!result.success);
373        assert_eq!(
374            result.output,
375            serde_json::Value::String("something went wrong".into())
376        );
377        assert!(result.metadata.is_empty());
378    }
379
380    #[test]
381    fn test_tool_result_metadata() {
382        let result = ToolResult::ok(serde_json::json!({"data": [1, 2, 3]}))
383            .with_metadata("source", serde_json::json!("database"))
384            .with_metadata("result_count", serde_json::json!(3))
385            .with_metadata("confidence", serde_json::json!(0.95));
386        assert!(result.success);
387        assert_eq!(result.metadata.len(), 3);
388        assert_eq!(result.metadata["source"], serde_json::json!("database"));
389        assert_eq!(result.metadata["result_count"], serde_json::json!(3));
390        assert_eq!(result.metadata["confidence"], serde_json::json!(0.95));
391    }
392
393    #[test]
394    fn test_tool_result_from_value() {
395        let value = serde_json::json!({"key": "val"});
396        let result: ToolResult = value.clone().into();
397        assert!(result.success);
398        assert_eq!(result.output, value);
399        assert!(result.metadata.is_empty());
400    }
401
402    #[tokio::test]
403    async fn test_execute_structured_default() {
404        let tool = FunctionTool::new(
405            "add",
406            "Add two numbers",
407            serde_json::json!({"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}),
408            |input| {
409                Box::pin(async move {
410                    let a = input["a"].as_f64().unwrap_or(0.0);
411                    let b = input["b"].as_f64().unwrap_or(0.0);
412                    Ok(serde_json::json!(a + b))
413                })
414            },
415        );
416
417        let result = tool
418            .execute_structured(serde_json::json!({"a": 3, "b": 4}))
419            .await
420            .unwrap();
421        assert!(result.success);
422        assert_eq!(result.output, serde_json::json!(7.0));
423        assert!(result.metadata.is_empty());
424    }
425
426    #[tokio::test]
427    async fn function_tool_propagates_error() {
428        let tool = FunctionTool::new(
429            "fail",
430            "Always fails",
431            serde_json::json!({"type": "object"}),
432            |_input| {
433                Box::pin(async move {
434                    Err(PeError::ToolExecution {
435                        tool: "fail".into(),
436                        reason: "intentional failure".into(),
437                    })
438                })
439            },
440        );
441
442        let result = tool.execute(serde_json::json!({})).await;
443        assert!(result.is_err());
444        let err = result.unwrap_err();
445        assert!(err.to_string().contains("intentional failure"));
446    }
447}
pe_tools/tool.rs

pe_tools/
tool.rs
