yapi_agent/
tool.rs

1//! Tools the agent can call.
2
3use std::sync::Arc;
4
5use futures_util::future::BoxFuture;
6use serde_json::{Map, Value};
7use tokio_util::sync::CancellationToken;
8use yapi_types::event::ToolResult;
9use yapi_types::message::ToolDeclaration;
10
11/// Whether a tool may run alongside other calls from the same response.
12#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
13pub enum ExecutionMode {
14    /// Runs concurrently with the other calls of its batch.
15    #[default]
16    Parallel,
17    /// Forces the whole batch to run one call at a time.
18    Sequential,
19}
20
21/// Receives partial results while a tool runs.
22pub type UpdateSink = Arc<dyn Fn(ToolResult) + Send + Sync>;
23
24/// A tool: a declaration for the model and an implementation.
25///
26/// `execute` returns `Err` with a message for failures, and the loop turns it
27/// into an error result. Cancellation is cooperative through `cancel`.
28pub trait Tool: Send + Sync {
29    /// Name, description and argument schema as sent to the model.
30    fn declaration(&self) -> &ToolDeclaration;
31
32    /// Human-readable name for the UI.
33    fn label(&self) -> &str {
34        &self.declaration().name
35    }
36
37    /// See [`ExecutionMode`].
38    fn execution_mode(&self) -> ExecutionMode {
39        ExecutionMode::Parallel
40    }
41
42    /// JSON Schema of the value the tool returns as `structuredContent`, for
43    /// tools whose results scripts consume (pi's `outputSchema`).
44    fn output_schema(&self) -> Option<&Value> {
45        None
46    }
47
48    /// Rewrites raw arguments before validation, for compatibility with older
49    /// argument shapes.
50    fn prepare_arguments(&self, arguments: Map<String, Value>) -> Map<String, Value> {
51        arguments
52    }
53
54    /// Runs the tool with validated arguments.
55    fn execute(
56        &self,
57        call_id: String,
58        args: Value,
59        cancel: CancellationToken,
60        updates: UpdateSink,
61    ) -> BoxFuture<'_, Result<ToolResult, String>>;
62}
yapi_agent/tool.rs

yapi_agent/
tool.rs
