use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use crate::id::{SessionId, TurnId};
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolDefinition {
pub name: String,
pub description: String,
pub parameters: serde_json::Value,
#[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
pub metadata: serde_json::Value,
}
impl ToolDefinition {
pub fn new(
name: impl Into<String>,
description: impl Into<String>,
parameters: serde_json::Value,
) -> Self {
Self {
name: name.into(),
description: description.into(),
parameters,
metadata: serde_json::Value::Null,
}
}
pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum DeferrablePolicy {
#[default]
Never,
Deferred,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub struct ToolPolicy {
pub deferrable: DeferrablePolicy,
}
impl ToolPolicy {
pub fn new(deferrable: DeferrablePolicy) -> Self {
Self { deferrable }
}
pub fn deferred() -> Self {
Self::new(DeferrablePolicy::Deferred)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolOutput {
pub content: String,
pub is_error: bool,
#[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
pub metadata: serde_json::Value,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub parts: Vec<crate::message::ContentPart>,
}
impl ToolOutput {
pub fn new(content: impl Into<String>, is_error: bool) -> Self {
Self {
content: content.into(),
is_error,
metadata: serde_json::Value::Null,
parts: Vec::new(),
}
}
pub fn text(content: impl Into<String>) -> Self {
Self::new(content, false)
}
pub fn error(content: impl Into<String>) -> Self {
Self::new(content, true)
}
pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
self.metadata = metadata;
self
}
pub fn with_parts(
mut self,
parts: impl IntoIterator<Item = crate::message::ContentPart>,
) -> Self {
self.parts = parts.into_iter().collect();
self
}
}
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolProgress {
pub message: String,
#[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
pub detail: serde_json::Value,
}
impl ToolProgress {
pub fn message(message: impl Into<String>) -> Self {
Self {
message: message.into(),
detail: serde_json::Value::Null,
}
}
pub fn with_detail(mut self, detail: serde_json::Value) -> Self {
self.detail = detail;
self
}
}
#[async_trait]
pub trait ToolProgressSink: Send + Sync {
async fn progress(&self, progress: ToolProgress);
}
#[derive(Clone, Default)]
#[non_exhaustive]
pub struct ToolContext {
pub session_id: SessionId,
pub turn_id: TurnId,
pub extensions: crate::extensions::Extensions,
pub cancellation: crate::cancellation::CancellationToken,
pub progress: Option<Arc<dyn ToolProgressSink>>,
}
impl std::fmt::Debug for ToolContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ToolContext")
.field("session_id", &self.session_id)
.field("turn_id", &self.turn_id)
.field("cancellation", &self.cancellation)
.field("progress", &self.progress.is_some())
.finish_non_exhaustive()
}
}
impl ToolContext {
pub fn new(session_id: SessionId, turn_id: TurnId) -> Self {
Self {
session_id,
turn_id,
extensions: crate::extensions::Extensions::new(),
cancellation: crate::cancellation::CancellationToken::new(),
progress: None,
}
}
pub fn with_extensions(mut self, extensions: crate::extensions::Extensions) -> Self {
self.extensions = extensions;
self
}
pub fn with_cancellation(
mut self,
cancellation: crate::cancellation::CancellationToken,
) -> Self {
self.cancellation = cancellation;
self
}
pub fn with_progress(mut self, progress: Arc<dyn ToolProgressSink>) -> Self {
self.progress = Some(progress);
self
}
pub async fn report_progress(&self, progress: ToolProgress) {
if let Some(sink) = &self.progress {
sink.progress(progress).await;
}
}
}
#[async_trait]
pub trait Tool: Send + Sync {
fn definition(&self) -> ToolDefinition;
async fn execute(&self, arguments: serde_json::Value, context: &ToolContext) -> ToolOutput;
fn policy(&self) -> ToolPolicy {
ToolPolicy::default()
}
}
type BoxedToolFuture = Pin<Box<dyn Future<Output = ToolOutput> + Send>>;
type BoxedToolFn = Box<dyn Fn(serde_json::Value) -> BoxedToolFuture + Send + Sync>;
pub struct FnTool {
definition: ToolDefinition,
handler: BoxedToolFn,
}
impl FnTool {
pub fn new<F, Fut>(
name: impl Into<String>,
description: impl Into<String>,
parameters: serde_json::Value,
handler: F,
) -> Self
where
F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
Fut: Future<Output = ToolOutput> + Send + 'static,
{
Self {
definition: ToolDefinition::new(name, description, parameters),
handler: Box::new(move |args| Box::pin(handler(args))),
}
}
}
#[async_trait]
impl Tool for FnTool {
fn definition(&self) -> ToolDefinition {
self.definition.clone()
}
async fn execute(&self, arguments: serde_json::Value, _context: &ToolContext) -> ToolOutput {
(self.handler)(arguments).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn metadata_hatch_round_trips_and_is_omitted_when_null() {
let plain = ToolDefinition::new("read", "Read a file.", json!({"type": "object"}));
let plain_json = serde_json::to_string(&plain).unwrap();
assert!(!plain_json.contains("metadata"));
let hinted = plain.with_metadata(json!({"hints": {"risk": "readonly"}}));
let back: ToolDefinition =
serde_json::from_str(&serde_json::to_string(&hinted).unwrap()).unwrap();
assert_eq!(back, hinted);
assert_eq!(back.metadata["hints"]["risk"], "readonly");
}
}