langchain_rust/agent/open_ai_tools/
builder.rs

1use std::sync::Arc;
2
3use crate::{
4    agent::AgentError,
5    chain::{options::ChainCallOptions, LLMChainBuilder},
6    language_models::{llm::LLM, options::CallOptions},
7    schemas::FunctionDefinition,
8    tools::Tool,
9};
10
11use super::{prompt::PREFIX, OpenAiToolAgent};
12
13pub struct OpenAiToolAgentBuilder {
14    tools: Option<Vec<Arc<dyn Tool>>>,
15    prefix: Option<String>,
16    options: Option<ChainCallOptions>,
17}
18
19impl OpenAiToolAgentBuilder {
20    pub fn new() -> Self {
21        Self {
22            tools: None,
23            prefix: None,
24            options: None,
25        }
26    }
27
28    pub fn tools(mut self, tools: &[Arc<dyn Tool>]) -> Self {
29        self.tools = Some(tools.to_vec());
30        self
31    }
32
33    pub fn prefix<S: Into<String>>(mut self, prefix: S) -> Self {
34        self.prefix = Some(prefix.into());
35        self
36    }
37
38    pub fn options(mut self, options: ChainCallOptions) -> Self {
39        self.options = Some(options);
40        self
41    }
42
43    pub fn build<L: LLM + 'static>(self, llm: L) -> Result<OpenAiToolAgent, AgentError> {
44        let tools = self.tools.unwrap_or_default();
45        let prefix = self.prefix.unwrap_or_else(|| PREFIX.to_string());
46        let mut llm = llm;
47
48        let prompt = OpenAiToolAgent::create_prompt(&prefix)?;
49        let default_options = ChainCallOptions::default().with_max_tokens(1000);
50        let functions = tools
51            .iter()
52            .map(FunctionDefinition::from_langchain_tool)
53            .collect::<Vec<FunctionDefinition>>();
54        llm.add_options(CallOptions::new().with_functions(functions));
55        let chain = Box::new(
56            LLMChainBuilder::new()
57                .prompt(prompt)
58                .llm(llm)
59                .options(self.options.unwrap_or(default_options))
60                .build()?,
61        );
62
63        Ok(OpenAiToolAgent { chain, tools })
64    }
65}