langchain_rust/agent/chat/
builder.rs1use 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}