ollama_rs/generation/tools/
mod.rs

1#[cfg_attr(docsrs, doc(cfg(feature = "tool-implementations")))]
2#[cfg(feature = "tool-implementations")]
3pub mod implementations;
4
5use std::{future::Future, pin::Pin};
6
7use schemars::{generate::SchemaSettings, JsonSchema, Schema};
8use serde::{de::DeserializeOwned, Deserialize, Serialize};
9use serde_json::Value;
10
11pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
12
13/// It's highly recommended that the `JsonSchema` has descriptions for all attributes.
14/// Descriptions can be defined with `#[schemars(description = "Hi I am an attribute")]` above each attribute
15// TODO enforce at compile-time
16pub trait Tool: Send + Sync {
17    type Params: Parameters;
18
19    fn name() -> &'static str;
20    fn description() -> &'static str;
21
22    /// Call the tool.
23    /// Note that returning an Err will cause it to be bubbled up. If you want the LLM to handle the error,
24    /// return that error as a string.
25    fn call(&mut self, parameters: Self::Params) -> impl Future<Output = Result<String>> + Send;
26}
27
28pub trait Parameters: DeserializeOwned + JsonSchema {}
29
30impl<P: DeserializeOwned + JsonSchema> Parameters for P {}
31
32pub(crate) trait ToolHolder: Send + Sync {
33    fn call(
34        &mut self,
35        parameters: Value,
36    ) -> Pin<Box<dyn Future<Output = Result<String>> + '_ + Send>>;
37}
38
39impl<T: Tool> ToolHolder for T {
40    fn call(
41        &mut self,
42        parameters: Value,
43    ) -> Pin<Box<dyn Future<Output = Result<String>> + '_ + Send>> {
44        Box::pin(async move {
45            // Json returned from the model can sometimes be in different formats, see https://github.com/pepperoni21/ollama-rs/issues/210
46            // This is a work-around for this issue.
47            let param_value = match serde_json::from_value(parameters.clone()) {
48                // We first try with the ToolCallFunction format
49                Ok(ToolCallFunction { name: _, arguments }) => arguments,
50                Err(_err) => match serde_json::from_value::<ToolInfo>(parameters.clone()) {
51                    Ok(ti) => ti.function.parameters.to_value(),
52                    Err(_err) => parameters,
53                },
54            };
55
56            let param = serde_json::from_value(param_value)?;
57
58            T::call(self, param).await
59        })
60    }
61}
62
63#[derive(Clone, Debug, Serialize, Deserialize)]
64pub struct ToolInfo {
65    #[serde(rename = "type")]
66    pub tool_type: ToolType,
67    pub function: ToolFunctionInfo,
68}
69
70impl ToolInfo {
71    /// Builds the JSON schema information Ollama needs to make a [`Tool`]
72    /// available to the model.
73    pub fn new<P: Parameters, T: Tool<Params = P>>() -> Self {
74        let mut settings = SchemaSettings::draft07();
75        settings.inline_subschemas = true;
76        let generator = settings.into_generator();
77
78        let parameters = generator.into_root_schema_for::<T::Params>();
79
80        Self {
81            tool_type: ToolType::Function,
82            function: ToolFunctionInfo {
83                name: T::name().to_string(),
84                description: T::description().to_string(),
85                parameters,
86            },
87        }
88    }
89}
90
91#[derive(Clone, Debug, Serialize, Deserialize)]
92pub enum ToolType {
93    #[serde(rename_all(deserialize = "PascalCase"))]
94    Function,
95}
96
97#[derive(Clone, Debug, Serialize, Deserialize)]
98pub struct ToolFunctionInfo {
99    pub name: String,
100    pub description: String,
101    pub parameters: Schema,
102}
103
104#[derive(Clone, Debug, Serialize, Deserialize)]
105pub struct ToolCall {
106    pub function: ToolCallFunction,
107}
108
109#[derive(Clone, Debug, Serialize, Deserialize)]
110pub struct ToolCallFunction {
111    pub name: String,
112    // I don't love this (the Value)
113    // But fixing it would be a big effort
114    // FIXME
115    #[serde(alias = "parameters")]
116    pub arguments: Value,
117}
ollama_rs/generation/tools/mod.rs

ollama_rs/generation/tools/
mod.rs
