rig_core/embeddings/
tool.rs

1//! The module defines the [ToolSchema] struct, which is used to embed an object that implements [crate::tool::PortableToolEmbedding]
2
3use crate::{Embed, tool::PortableToolEmbedding};
4use serde::Serialize;
5
6use super::embed::EmbedError;
7
8/// Embeddable document that is used as an intermediate representation of a tool when
9/// RAGging tools.
10#[derive(Clone, Serialize, Default, Eq, PartialEq)]
11pub struct ToolSchema {
12    pub name: String,
13    pub context: serde_json::Value,
14    pub embedding_docs: Vec<String>,
15}
16
17impl Embed for ToolSchema {
18    fn embed(&self, embedder: &mut super::embed::TextEmbedder) -> Result<(), EmbedError> {
19        for doc in &self.embedding_docs {
20            embedder.embed(doc.clone());
21        }
22        Ok(())
23    }
24}
25
26impl ToolSchema {
27    /// Convert an embedding-backed tool to a [`ToolSchema`].
28    ///
29    /// # Example
30    /// ```rust
31    /// use rig_core::{
32    ///     embeddings::ToolSchema,
33    ///     tool::{PortableTool, PortableToolEmbedding},
34    /// };
35    ///
36    /// #[derive(Debug, thiserror::Error)]
37    /// #[error("Nothing error")]
38    /// struct NothingError;
39    ///
40    /// #[derive(Debug, thiserror::Error)]
41    /// #[error("Init error")]
42    /// struct InitError;
43    ///
44    /// struct Nothing;
45    /// impl PortableTool for Nothing {
46    ///     const NAME: &'static str = "nothing";
47    ///
48    ///     type Args = ();
49    ///     type Output = ();
50    ///     type Error = NothingError;
51    ///
52    ///     fn description(&self) -> String {
53    ///         "nothing".to_string()
54    ///     }
55    ///
56    ///     fn parameters(&self) -> serde_json::Value {
57    ///         serde_json::json!({})
58    ///     }
59    ///
60    ///     async fn call(&self, _args: Self::Args) -> Result<Self::Output, Self::Error> {
61    ///         Ok(())
62    ///     }
63    /// }
64    ///
65    /// impl PortableToolEmbedding for Nothing {
66    ///     type InitError = InitError;
67    ///     type Context = ();
68    ///     type State = ();
69    ///
70    ///     fn init(_state: Self::State, _context: Self::Context) -> Result<Self, Self::InitError> {
71    ///         Ok(Nothing)
72    ///     }
73    ///
74    ///     fn embedding_docs(&self) -> Vec<String> {
75    ///         vec!["Do nothing.".into()]
76    ///     }
77    ///
78    ///     fn context(&self) -> Self::Context {}
79    /// }
80    ///
81    /// let tool = ToolSchema::try_from(&Nothing).unwrap();
82    ///
83    /// assert_eq!(tool.name, "nothing".to_string());
84    /// assert_eq!(tool.embedding_docs, vec!["Do nothing.".to_string()]);
85    /// ```
86    pub fn try_from<T>(tool: &T) -> Result<Self, EmbedError>
87    where
88        T: PortableToolEmbedding + 'static,
89    {
90        Ok(ToolSchema {
91            name: T::NAME.to_string(),
92            context: serde_json::to_value(tool.context()).map_err(EmbedError::new)?,
93            embedding_docs: tool.embedding_docs(),
94        })
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use std::convert::Infallible;
101
102    use super::ToolSchema;
103    use crate::tool::{PortableTool, PortableToolEmbedding, ToolExecutionError};
104
105    struct NamedTool;
106
107    impl PortableTool for NamedTool {
108        const NAME: &'static str = "static_name";
109
110        type Error = rig::tool::ToolExecutionError;
111
112        type Args = ();
113        type Output = ();
114
115        fn description(&self) -> String {
116            "A statically named tool".to_string()
117        }
118
119        fn parameters(&self) -> serde_json::Value {
120            serde_json::json!({})
121        }
122
123        async fn call(&self, _args: Self::Args) -> Result<Self::Output, ToolExecutionError> {
124            Ok(())
125        }
126    }
127
128    impl PortableToolEmbedding for NamedTool {
129        type InitError = Infallible;
130        type Context = ();
131        type State = ();
132
133        fn embedding_docs(&self) -> Vec<String> {
134            vec!["named tool".to_string()]
135        }
136
137        fn context(&self) -> Self::Context {}
138
139        fn init(_state: Self::State, _context: Self::Context) -> Result<Self, Self::InitError> {
140            Ok(Self)
141        }
142    }
143
144    #[test]
145    fn try_from_uses_canonical_tool_name() {
146        let schema = ToolSchema::try_from(&NamedTool).unwrap();
147
148        assert_eq!(schema.name, NamedTool::NAME);
149    }
150}
rig_core/embeddings/tool.rs

rig_core/embeddings/
tool.rs
