llm_chain/tools/
collection.rs

1use super::tool::{Tool, ToolError};
2use crate::parsing::{find_yaml, ExtractionError};
3use crate::prompt::StringTemplate;
4use serde::{Deserialize, Serialize};
5use thiserror::Error;
6
7#[derive(Default)]
8pub struct ToolCollection<T> {
9    tools: Vec<T>,
10}
11
12#[derive(Error, Debug)]
13pub enum ToolUseError<E: ToolError> {
14    #[error("Model is not trying to invoke tools")]
15    NoToolInvocation,
16    #[error("You must output at most one tool invocation")]
17    MultipleInvocations,
18    #[error("Tool not found")]
19    ToolNotFound,
20    #[error("You must output YAML: {0}")]
21    InvalidYaml(#[from] ExtractionError),
22    #[error("Invalid format: {0}")]
23    InvalidFormat(#[from] serde_yaml::Error),
24    #[error("Tool invocation failed: {0}")]
25    ToolInvocationFailed(String),
26    #[error(transparent)]
27    ToolError(#[from] E),
28}
29
30impl<T> ToolCollection<T>
31where
32    T: Tool + Send + Sync,
33{
34    pub fn new() -> Self {
35        Self { tools: vec![] }
36    }
37
38    pub fn add_tool(&mut self, tool: T) {
39        self.tools.push(tool);
40    }
41
42    pub async fn invoke(
43        &self,
44        name: &str,
45        input: &serde_yaml::Value,
46    ) -> Result<serde_yaml::Value, ToolUseError<<T as Tool>::Error>> {
47        let tool = self
48            .tools
49            .iter()
50            .find(|t| t.matches(name))
51            .ok_or(ToolUseError::ToolNotFound)?;
52        tool.invoke(input.clone()).await.map_err(|e| e.into())
53    }
54
55    pub fn get_tool_invocation(
56        &self,
57        data: &str,
58    ) -> Result<ToolInvocationInput, ToolUseError<<T as Tool>::Error>> {
59        let tool_invocations: Vec<ToolInvocationInput> = find_yaml::<ToolInvocationInput>(data)?;
60        if tool_invocations.len() > 1 {
61            return Err(ToolUseError::MultipleInvocations);
62        }
63        tool_invocations
64            .first()
65            .cloned()
66            .ok_or(ToolUseError::NoToolInvocation)
67    }
68
69    /// Process chat input and execute the appropriate tool.
70    ///
71    /// The input string should contain a YAML block describing the tool invocation.
72    /// The YAML block should have a `command` field and an `input` field.
73    ///
74    /// # Errors
75    ///
76    /// Returns an `OpaqueError` variant if the input is not a valid YAML or
77    /// if the specified tool is not found.
78    pub async fn process_chat_input(
79        &self,
80        data: &str,
81    ) -> Result<String, ToolUseError<<T as Tool>::Error>> {
82        let tool_invocation = self.get_tool_invocation(data)?;
83        let output = self
84            .invoke(&tool_invocation.command, &tool_invocation.input)
85            .await?;
86        serde_yaml::to_string(&output).map_err(|e| e.into())
87    }
88
89    /// Generate a YAML-formatted string describing the available tools.
90    pub fn describe(&self) -> Result<String, ToolUseError<<T as Tool>::Error>> {
91        let des: Vec<_> = self.tools.iter().map(|t| t.description()).collect();
92        serde_yaml::to_string(&des).map_err(|e| e.into())
93    }
94
95    /// Generate a prompt template for the tool collection. Combine it with a normal prompt template to perform your task.
96    pub fn to_prompt_template(&self) -> Result<StringTemplate, ToolUseError<<T as Tool>::Error>> {
97        Ok(StringTemplate::combine(vec![
98            StringTemplate::static_string(include_str!("./tool_prompt_prefix.txt").to_string()),
99            StringTemplate::static_string(self.describe()?),
100            StringTemplate::static_string("\n\n"),
101        ]))
102    }
103}
104
105#[derive(Serialize, Deserialize, Debug, Clone)]
106pub struct ToolInvocationInput {
107    pub command: String,
108    pub input: serde_yaml::Value,
109}
llm_chain/tools/collection.rs

llm_chain/tools/
collection.rs
