use crate::runnables::{LcelError, Runnable, RunnableConfig};
use async_trait::async_trait;
use schemars::JsonSchema;
use serde::{de::DeserializeOwned, Serialize};
use serde_json::Value;
use std::sync::Arc;
#[async_trait]
pub trait BaseTool: Send + Sync {
fn name(&self) -> &str;
fn description(&self) -> &str;
async fn run(&self, input: String) -> Result<String, ToolError>;
fn args_schema(&self) -> Option<Value> {
None
}
fn return_direct(&self) -> bool {
false
}
async fn handle_error(&self, error: ToolError) -> String {
format!("Tool '{}' execution failed: {}", self.name(), error)
}
}
#[async_trait]
pub trait Tool: Send + Sync {
type Input: DeserializeOwned + JsonSchema + Send + Sync + 'static;
type Output: Serialize + Send + Sync;
async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError>;
fn args_schema(&self) -> Option<Value> {
use schemars::schema_for;
serde_json::to_value(schema_for!(Self::Input)).ok()
}
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum ToolError {
#[error("Invalid input: {0}")]
InvalidInput(String),
#[error("Execution failed: {0}")]
ExecutionFailed(String),
#[error("Timeout: {0} seconds")]
Timeout(u64),
#[error("Tool not found: {0}")]
ToolNotFound(String),
#[error("MCP error [{code}]: {message}")]
McpError {
code: i32,
message: String,
data: Option<Value>,
},
}
use super::ToolDefinition;
pub fn to_tool_definition(tool: &dyn BaseTool) -> ToolDefinition {
ToolDefinition::new(tool.name(), tool.description()).with_parameters(
tool.args_schema()
.unwrap_or(serde_json::json!({"type": "object"})),
)
}
#[async_trait]
impl Runnable<String, String> for Arc<dyn BaseTool> {
type Error = LcelError;
async fn invoke(
&self,
input: String,
_config: Option<RunnableConfig>,
) -> Result<String, LcelError> {
BaseTool::run(&**self, input).await.map_err(LcelError::from)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
struct EchoTool;
#[async_trait]
impl BaseTool for EchoTool {
fn name(&self) -> &str {
"echo"
}
fn description(&self) -> &str {
"回显输入"
}
async fn run(&self, input: String) -> Result<String, ToolError> {
Ok(format!("echo: {input}"))
}
}
#[tokio::test]
async fn arc_tool_is_runnable() {
let tool: Arc<dyn BaseTool> = Arc::new(EchoTool);
let result = tool.invoke("hi".to_string(), None).await.unwrap();
assert_eq!(result, "echo: hi");
}
#[tokio::test]
async fn arc_tool_pipes() {
use crate::runnables::{RunnableExt, RunnableLambda};
let tool: Arc<dyn BaseTool> = Arc::new(EchoTool);
let chain = tool.pipe(RunnableLambda::new_sync(|s: String| s.to_uppercase()));
let result = chain.invoke("hi".to_string(), None).await.unwrap();
assert_eq!(result, "ECHO: HI");
}
#[tokio::test]
async fn arc_tool_error_maps_to_lcel() {
struct FailingTool;
#[async_trait]
impl BaseTool for FailingTool {
fn name(&self) -> &str {
"fail"
}
fn description(&self) -> &str {
"总是失败"
}
async fn run(&self, _input: String) -> Result<String, ToolError> {
Err(ToolError::ExecutionFailed("boom".to_string()))
}
}
let tool: Arc<dyn BaseTool> = Arc::new(FailingTool);
let err = tool.invoke("x".to_string(), None).await.unwrap_err();
assert!(matches!(err, LcelError::Tool(ref msg) if msg.contains("boom")));
}
}