Skip to main content

langchain_rust/tools/function_tool/
mod.rs

1use std::{error::Error, future::Future, marker::PhantomData};
2
3use async_trait::async_trait;
4use serde::de::DeserializeOwned;
5use serde_json::{json, Value};
6
7use super::Tool;
8
9/// A [`Tool`] built from an async function, without writing a struct and a trait impl.
10///
11/// The function takes any `serde` type: the arguments the model sends are deserialized
12/// into it. `parameters` is the JSON schema advertised to the model; it should describe
13/// that type.
14///
15/// ```rust,ignore
16/// #[derive(serde::Deserialize)]
17/// struct AddInput { a: f64, b: f64 }
18///
19/// let adder = FunctionTool::new(
20///     "adder",
21///     "Adds two numbers together",
22///     json!({
23///         "type": "object",
24///         "properties": { "a": {"type": "number"}, "b": {"type": "number"} },
25///         "required": ["a", "b"]
26///     }),
27///     |input: AddInput| async move {
28///         Ok::<_, String>(format!("The sum of {} and {} is {}", input.a, input.b, input.a + input.b))
29///     },
30/// );
31/// let tools: Vec<Arc<dyn Tool>> = vec![Arc::new(adder)];
32/// ```
33///
34/// For a tool that takes a single string, [`FunctionTool::text`] fills in the schema.
35///
36/// Input that doesn't deserialize comes back to the model as a tool error naming the
37/// expected schema, so it can retry; the function isn't called.
38pub struct FunctionTool<T, F> {
39    name: String,
40    description: String,
41    parameters: Value,
42    func: F,
43    // `fn() -> T` keeps the tool Send + Sync whatever `T` is; no `T` is ever stored.
44    _input: PhantomData<fn() -> T>,
45}
46
47impl<T, F> FunctionTool<T, F> {
48    pub fn new<N: Into<String>, D: Into<String>>(
49        name: N,
50        description: D,
51        parameters: Value,
52        func: F,
53    ) -> Self {
54        Self {
55            name: name.into(),
56            description: description.into(),
57            parameters,
58            func,
59            _input: PhantomData,
60        }
61    }
62}
63
64impl<F> FunctionTool<String, F> {
65    /// A tool whose function takes the model's input as one string. The advertised schema
66    /// is `{"input": string}`, described by `description`.
67    pub fn text<N: Into<String>, D: Into<String>>(name: N, description: D, func: F) -> Self {
68        let description = description.into();
69        let parameters = json!({
70            "type": "object",
71            "properties": {
72                "input": { "type": "string", "description": description }
73            },
74            "required": ["input"]
75        });
76        Self::new(name, description, parameters, func)
77    }
78}
79
80impl<T: DeserializeOwned, F> FunctionTool<T, F> {
81    /// Deserializes the model's arguments. Besides the value itself, accepts the
82    /// `{"input": ...}` wrapper the default schema asks for, arguments double-encoded as a
83    /// JSON string, and a bare value for a string input (e.g. `42` when the model skipped
84    /// the quotes).
85    fn decode(&self, value: Value) -> Result<T, Box<dyn Error>> {
86        let first_err = match serde_json::from_value::<T>(value.clone()) {
87            Ok(input) => return Ok(input),
88            Err(e) => e,
89        };
90        let fallbacks = [
91            value.get("input").cloned(),
92            value.as_str().and_then(|s| serde_json::from_str(s).ok()),
93            (!value.is_string()).then(|| Value::String(value.to_string())),
94        ];
95        fallbacks
96            .into_iter()
97            .flatten()
98            .find_map(|candidate| serde_json::from_value::<T>(candidate).ok())
99            .ok_or_else(|| {
100                format!(
101                    "invalid input for tool `{}`: {first_err}. Expected arguments matching: {}",
102                    self.name, self.parameters
103                )
104                .into()
105            })
106    }
107}
108
109#[async_trait]
110impl<T, F, Fut, O, E> Tool for FunctionTool<T, F>
111where
112    T: DeserializeOwned + Send + 'static,
113    F: Fn(T) -> Fut + Send + Sync,
114    Fut: Future<Output = Result<O, E>> + Send,
115    O: ToString,
116    E: Into<Box<dyn Error>>,
117{
118    fn name(&self) -> String {
119        self.name.clone()
120    }
121
122    fn description(&self) -> String {
123        self.description.clone()
124    }
125
126    fn parameters(&self) -> Value {
127        self.parameters.clone()
128    }
129
130    /// Keeps the arguments as JSON (the default turns objects into a string); anything
131    /// that isn't JSON is passed on as a string.
132    async fn parse_input(&self, input: &str) -> Value {
133        serde_json::from_str(input).unwrap_or_else(|_| Value::String(input.to_string()))
134    }
135
136    async fn run(&self, input: Value) -> Result<String, Box<dyn Error>> {
137        let input = self.decode(input)?;
138        (self.func)(input)
139            .await
140            .map(|output| output.to_string())
141            .map_err(Into::into)
142    }
143}
144
145#[cfg(test)]
146mod tests;