use std::sync::Arc;
use serde::{Deserialize, Serialize};
use crate::{
completion::ToolDefinition,
wasm_compat::{WasmBoxedFuture, WasmCompatSend, WasmCompatSync},
};
use super::{IntoToolOutput, ToolExecutionError, ToolOutput};
pub trait PortableTool: Sized + WasmCompatSend + WasmCompatSync {
const NAME: &'static str;
type Args: for<'de> Deserialize<'de> + WasmCompatSend + WasmCompatSync;
type Output: IntoToolOutput + WasmCompatSend;
type Error: std::error::Error + WasmCompatSend + WasmCompatSync + 'static;
fn description(&self) -> String;
fn parameters(&self) -> serde_json::Value;
fn map_error(&self, error: Self::Error) -> ToolExecutionError {
ToolExecutionError::from_error(error)
}
fn call(
&self,
arguments: Self::Args,
) -> impl Future<Output = Result<Self::Output, Self::Error>> + WasmCompatSend;
}
pub trait PortableToolEmbedding: PortableTool {
type InitError: std::error::Error + WasmCompatSend + WasmCompatSync + 'static;
type Context: for<'de> Deserialize<'de> + Serialize;
type State: WasmCompatSend;
fn embedding_docs(&self) -> Vec<String>;
fn context(&self) -> Self::Context;
fn init(state: Self::State, context: Self::Context) -> Result<Self, Self::InitError>;
}
trait PortableDynamicCallback:
Fn(serde_json::Value) -> WasmBoxedFuture<'static, Result<ToolOutput, ToolExecutionError>>
+ WasmCompatSend
+ WasmCompatSync
{
}
impl<F> PortableDynamicCallback for F where
F: Fn(serde_json::Value) -> WasmBoxedFuture<'static, Result<ToolOutput, ToolExecutionError>>
+ WasmCompatSend
+ WasmCompatSync
{
}
#[derive(Clone)]
pub struct PortableDynamicTool {
name: String,
description: String,
parameters: serde_json::Value,
callback: Arc<dyn PortableDynamicCallback>,
}
impl std::fmt::Debug for PortableDynamicTool {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("PortableDynamicTool")
.field("name", &self.name)
.field("description", &self.description)
.field("parameters", &self.parameters)
.finish_non_exhaustive()
}
}
impl PortableDynamicTool {
pub fn new<F>(
name: impl Into<String>,
description: impl Into<String>,
parameters: serde_json::Value,
callback: F,
) -> Self
where
F: Fn(
serde_json::Value,
) -> WasmBoxedFuture<'static, Result<ToolOutput, ToolExecutionError>>
+ WasmCompatSend
+ WasmCompatSync
+ 'static,
{
Self {
name: name.into(),
description: description.into(),
parameters,
callback: Arc::new(callback),
}
}
pub fn name(&self) -> &str {
&self.name
}
pub fn definition(&self) -> ToolDefinition {
ToolDefinition {
name: self.name.clone(),
description: self.description.clone(),
parameters: self.parameters.clone(),
}
}
pub async fn execute(
&self,
arguments: serde_json::Value,
) -> Result<ToolOutput, ToolExecutionError> {
(self.callback)(arguments).await
}
}
pub fn portable_tool_definition<T>(tool: &T) -> ToolDefinition
where
T: PortableTool,
{
ToolDefinition {
name: T::NAME.to_owned(),
description: tool.description(),
parameters: tool.parameters(),
}
}
#[cfg(test)]
mod tests {
use std::convert::Infallible;
use serde::{Deserialize, Serialize};
use super::*;
#[derive(Deserialize)]
struct AddArgs {
left: i64,
right: i64,
}
#[derive(Serialize)]
struct Sum {
value: i64,
}
struct Add;
impl PortableTool for Add {
const NAME: &'static str = "add";
type Args = AddArgs;
type Output = Sum;
type Error = Infallible;
fn description(&self) -> String {
"Add two integers".to_string()
}
fn parameters(&self) -> serde_json::Value {
serde_json::json!({"type": "object"})
}
async fn call(&self, arguments: Self::Args) -> Result<Self::Output, Self::Error> {
Ok(Sum {
value: arguments.left + arguments.right,
})
}
}
#[tokio::test]
async fn portable_tools_execute_without_runtime_context() {
let output = Add.call(AddArgs { left: 2, right: 3 }).await;
let Ok(output) = output;
assert_eq!(output.value, 5);
assert_eq!(portable_tool_definition(&Add).name, "add");
}
#[tokio::test]
async fn portable_dynamic_tools_receive_owned_arguments() {
let tool = PortableDynamicTool::new(
"echo",
"Echo a JSON value",
serde_json::json!({"type": "object"}),
|arguments| Box::pin(async move { Ok(ToolOutput::json(arguments)) }),
);
let arguments = serde_json::json!({"value": "hello"});
let output = tool.execute(arguments.clone()).await;
assert!(output.is_ok());
let Ok(output) = output else {
return;
};
assert_eq!(output.as_json(), Some(&arguments));
}
}