langchain_rust/chain/combine_documents/
mod.rs

1//! Chains that combine many documents into one answer, beyond
2//! [`StuffDocument`](crate::chain::StuffDocument): refine, map-reduce, map-rerank,
3//! summarization, and [`RetrievalQA`] on top of a retriever.
4mod map_reduce;
5mod map_rerank;
6mod prompts;
7mod refine;
8mod retrieval_qa;
9
10pub use map_reduce::MapReduceDocuments;
11pub use map_rerank::{
12    MapRerankDocuments, RankedAnswer, MAP_RERANK_INTERMEDIATE_STEPS_KEY, MAP_RERANK_SCORE_KEY,
13};
14pub use refine::RefineDocuments;
15pub use retrieval_qa::{RetrievalQA, RETRIEVAL_QA_SOURCE_DOCUMENTS_KEY};
16
17pub(crate) use prompts::STUFF_QA_TEMPLATE;
18
19use crate::{
20    chain::{LLMChain, LLMChainBuilder, StuffDocument},
21    language_models::llm::LLM,
22    prompt::{PromptTemplate, TemplateFormat},
23};
24
25fn llm_chain(llm: Box<dyn LLM>, template: &str, variables: &[&str]) -> LLMChain {
26    let prompt = PromptTemplate::new(
27        template.to_string(),
28        variables.iter().map(|v| v.to_string()).collect(),
29        TemplateFormat::Jinja2,
30    );
31    LLMChainBuilder::new()
32        .prompt(prompt)
33        .llm(llm)
34        .build()
35        .expect("prompt and llm are set")
36}
37
38impl StuffDocument {
39    /// Summarizes all documents in one prompt. Input: `input_documents`.
40    pub fn load_summarize<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
41        StuffDocument::new(llm_chain(
42            llm.into(),
43            prompts::SUMMARIZE_TEMPLATE,
44            &["context"],
45        ))
46    }
47}
48
49impl RefineDocuments {
50    /// Answers `question` from the first document, then refines the answer with each of the
51    /// others. Inputs: `input_documents`, `question`.
52    pub fn load_qa<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
53        let llm = llm.into();
54        RefineDocuments::new(
55            llm_chain(
56                llm.clone_box(),
57                prompts::STUFF_QA_TEMPLATE,
58                &["context", "question"],
59            ),
60            llm_chain(
61                llm,
62                prompts::REFINE_QA_TEMPLATE,
63                &["question", "existing_answer", "context"],
64            ),
65        )
66    }
67
68    /// Summarizes the first document, then refines the summary with each of the others.
69    /// Input: `input_documents`.
70    pub fn load_summarize<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
71        let llm = llm.into();
72        RefineDocuments::new(
73            llm_chain(llm.clone_box(), prompts::SUMMARIZE_TEMPLATE, &["context"]),
74            llm_chain(
75                llm,
76                prompts::REFINE_SUMMARIZE_TEMPLATE,
77                &["existing_answer", "context"],
78            ),
79        )
80    }
81}
82
83impl MapReduceDocuments {
84    /// Extracts the passages relevant to `question` from each document, then answers from
85    /// those. Inputs: `input_documents`, `question`.
86    pub fn load_qa<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
87        let llm = llm.into();
88        MapReduceDocuments::new(
89            llm_chain(
90                llm.clone_box(),
91                prompts::MAP_REDUCE_MAP_QA_TEMPLATE,
92                &["context", "question"],
93            ),
94            StuffDocument::new(llm_chain(
95                llm,
96                prompts::MAP_REDUCE_COMBINE_QA_TEMPLATE,
97                &["context", "question"],
98            )),
99        )
100    }
101
102    /// Summarizes each document, then summarizes the summaries. Input: `input_documents`.
103    pub fn load_summarize<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
104        let llm = llm.into();
105        MapReduceDocuments::new(
106            llm_chain(llm.clone_box(), prompts::SUMMARIZE_TEMPLATE, &["context"]),
107            StuffDocument::load_summarize(llm),
108        )
109    }
110}
111
112impl MapRerankDocuments {
113    /// Answers `question` from each document with a self-reported score and returns the
114    /// best-scored answer. Inputs: `input_documents`, `question`.
115    pub fn load_qa<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
116        MapRerankDocuments::new(llm_chain(
117            llm.into(),
118            prompts::MAP_RERANK_QA_TEMPLATE,
119            &["context", "question"],
120        ))
121    }
122}
123
124#[cfg(test)]
125mod tests;
langchain_rust/chain/combine_documents/mod.rs

langchain_rust/chain/combine_documents/
mod.rs
