1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

use crate::{
    chain::{options::ChainCallOptions, ChainError, LLMChainBuilder},
    language_models::llm::LLM,
    output_parsers::OutputParser,
    prompt::FormatPrompter,
    template_jinja2,
};

use super::StuffDocument;

pub struct StuffDocumentBuilder {
    llm: Option<Box<dyn LLM>>,
    options: Option<ChainCallOptions>,
    output_key: Option<String>,
    output_parser: Option<Box<dyn OutputParser>>,
    prompt: Option<Box<dyn FormatPrompter>>,
}
impl StuffDocumentBuilder {
    pub fn new() -> Self {
        Self {
            llm: None,
            options: None,
            output_key: None,
            output_parser: None,
            prompt: None,
        }
    }

    pub fn llm<L: Into<Box<dyn LLM>>>(mut self, llm: L) -> Self {
        self.llm = Some(llm.into());
        self
    }

    pub fn options(mut self, options: ChainCallOptions) -> Self {
        self.options = Some(options);
        self
    }

    pub fn output_key<S: Into<String>>(mut self, output_key: S) -> Self {
        self.output_key = Some(output_key.into());
        self
    }

    ///If you want to add a custom prompt,keep in mind which variables are obligatory.
    pub fn prompt<P: Into<Box<dyn FormatPrompter>>>(mut self, prompt: P) -> Self {
        self.prompt = Some(prompt.into());
        self
    }

    pub fn build(self) -> Result<StuffDocument, ChainError> {
        let llm = self
            .llm
            .ok_or_else(|| ChainError::MissingObject("LLM must be set".into()))?;
        let prompt = match self.prompt {
            Some(prompt) => prompt,
            None => Box::new(template_jinja2!(
                DEFAULT_STUFF_QA_TEMPLATE,
                "context",
                "question"
            )),
        };

        let llm_chain = {
            let mut builder = LLMChainBuilder::new()
                .prompt(prompt)
                .options(self.options.unwrap_or_default())
                .llm(llm);
            if let Some(output_parser) = self.output_parser {
                builder = builder.output_parser(output_parser);
            }

            builder.build()?
        };

        Ok(StuffDocument::new(llm_chain))
    }
}

const 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.

{{context}}

Question:{{question}}
Helpful Answer:
"#;