langchain_rust/chain/stuff_documents/
builder.rs

1use crate::{
2    chain::{options::ChainCallOptions, ChainError, LLMChainBuilder},
3    language_models::llm::LLM,
4    output_parsers::OutputParser,
5    prompt::FormatPrompter,
6    template_jinja2,
7};
8
9use super::StuffDocument;
10
11pub struct StuffDocumentBuilder {
12    llm: Option<Box<dyn LLM>>,
13    options: Option<ChainCallOptions>,
14    output_key: Option<String>,
15    output_parser: Option<Box<dyn OutputParser>>,
16    prompt: Option<Box<dyn FormatPrompter>>,
17}
18impl StuffDocumentBuilder {
19    pub fn new() -> Self {
20        Self {
21            llm: None,
22            options: None,
23            output_key: None,
24            output_parser: None,
25            prompt: None,
26        }
27    }
28
29    pub fn llm<L: Into<Box<dyn LLM>>>(mut self, llm: L) -> Self {
30        self.llm = Some(llm.into());
31        self
32    }
33
34    pub fn options(mut self, options: ChainCallOptions) -> Self {
35        self.options = Some(options);
36        self
37    }
38
39    pub fn output_key<S: Into<String>>(mut self, output_key: S) -> Self {
40        self.output_key = Some(output_key.into());
41        self
42    }
43
44    ///If you want to add a custom prompt,keep in mind which variables are obligatory.
45    pub fn prompt<P: Into<Box<dyn FormatPrompter>>>(mut self, prompt: P) -> Self {
46        self.prompt = Some(prompt.into());
47        self
48    }
49
50    pub fn build(self) -> Result<StuffDocument, ChainError> {
51        let llm = self
52            .llm
53            .ok_or_else(|| ChainError::MissingObject("LLM must be set".into()))?;
54        let prompt = match self.prompt {
55            Some(prompt) => prompt,
56            None => Box::new(template_jinja2!(
57                DEFAULT_STUFF_QA_TEMPLATE,
58                "context",
59                "question"
60            )),
61        };
62
63        let llm_chain = {
64            let mut builder = LLMChainBuilder::new()
65                .prompt(prompt)
66                .options(self.options.unwrap_or_default())
67                .llm(llm);
68            if let Some(output_parser) = self.output_parser {
69                builder = builder.output_parser(output_parser);
70            }
71
72            builder.build()?
73        };
74
75        Ok(StuffDocument::new(llm_chain))
76    }
77}
78
79const DEFAULT_STUFF_QA_TEMPLATE: &str = r#"Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
80
81{{context}}
82
83Question:{{question}}
84Helpful Answer:
85"#;
langchain_rust/chain/stuff_documents/builder.rs

langchain_rust/chain/stuff_documents/
builder.rs
