langchain_rust/agent/chat/
builder.rs

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