lc_chains/document_chains/
map_reduce.rs

1// lc-chains/src/document_chains/map_reduce.rs
2//! MapReduceDocumentsChain - processes documents in parallel then merges results.
3
4use async_trait::async_trait;
5use futures_util::future::try_join_all;
6use lc_core::language_models::LLMResult;
7use lc_core::{BaseChatModel, Runnable};
8use lc_schema::Message;
9use lc_shared::document::Document;
10use serde_json::Value;
11use std::collections::HashMap;
12
13use crate::base::{BaseChain, ChainError, ChainResult};
14
15/// Default Map processing prompt template.
16pub(crate) const DEFAULT_MAP_PROMPT: &str = "Answer the user's question based on the following document content. Provide a concise answer based on the document content.
17
18Document content:
19{context}
20
21Question: {input}
22
23Answer based on this document:";
24
25/// Default Reduce merge prompt template.
26pub(crate) const DEFAULT_REDUCE_PROMPT: &str = "Below are answers from multiple documents. Please merge them into a single complete and coherent final answer.
27
28Answers from each document:
29{summaries}
30
31Original question: {input}
32
33Final consolidated answer:";
34
35/// MapReduceDocumentsChain
36///
37/// Processes documents in two steps:
38/// 1. Map: Calls LLM independently for each document to generate an answer
39/// 2. Reduce: Merges all independent answers into a final answer
40pub struct MapReduceDocumentsChain<M: BaseChatModel> {
41    llm: M,
42    map_prompt_template: String,
43    reduce_prompt_template: String,
44    document_variable_name: String,
45    input_key: String,
46    output_key: String,
47    name: String,
48    verbose: bool,
49}
50
51impl<M: BaseChatModel> MapReduceDocumentsChain<M> {
52    pub fn new(llm: M) -> Self {
53        Self {
54            llm,
55            map_prompt_template: DEFAULT_MAP_PROMPT.to_string(),
56            reduce_prompt_template: DEFAULT_REDUCE_PROMPT.to_string(),
57            document_variable_name: "context".to_string(),
58            input_key: "input".to_string(),
59            output_key: "output".to_string(),
60            name: "map_reduce_documents".to_string(),
61            verbose: false,
62        }
63    }
64
65    pub fn with_map_prompt(mut self, template: impl Into<String>) -> Self {
66        self.map_prompt_template = template.into();
67        self
68    }
69
70    pub fn with_reduce_prompt(mut self, template: impl Into<String>) -> Self {
71        self.reduce_prompt_template = template.into();
72        self
73    }
74
75    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
76        self.document_variable_name = name.into();
77        self
78    }
79
80    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
81        self.input_key = key.into();
82        self
83    }
84
85    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
86        self.output_key = key.into();
87        self
88    }
89
90    pub fn with_name(mut self, name: impl Into<String>) -> Self {
91        self.name = name.into();
92        self
93    }
94
95    pub fn with_verbose(mut self, verbose: bool) -> Self {
96        self.verbose = verbose;
97        self
98    }
99
100    pub fn build_map_prompt(&self, context: &str, input: &str) -> String {
101        self.map_prompt_template
102            .replace(&format!("{{{}}}", self.document_variable_name), context)
103            .replace("{input}", input)
104    }
105
106    pub fn build_reduce_prompt(&self, summaries: &[String], input: &str) -> String {
107        let summaries_text = summaries
108            .iter()
109            .enumerate()
110            .map(|(i, s)| format!("Answer from document {}:\n{}", i + 1, s))
111            .collect::<Vec<_>>()
112            .join("\n\n");
113
114        self.reduce_prompt_template
115            .replace("{summaries}", &summaries_text)
116            .replace("{input}", input)
117    }
118
119    /// Map phase: call LLM for a single document.
120    async fn map_document(
121        &self,
122        doc: &Document,
123        input: &str,
124        index: usize,
125    ) -> Result<String, ChainError>
126    where
127        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
128    {
129        let prompt = self.build_map_prompt(&doc.content, input);
130
131        if self.verbose {
132            println!("\n--- Map document {} ---", index + 1);
133        }
134
135        let messages = vec![Message::human(&prompt)];
136        let response = self.llm.invoke(messages, None).await.map_err(|e| {
137            ChainError::ExecutionError(format!("Map call failed (document {}): {}", index + 1, e))
138        })?;
139
140        if self.verbose {
141            println!("Document {} answer: {}", index + 1, response.content);
142        }
143
144        Ok(response.content)
145    }
146
147    /// Invoke with documents and input directly.
148    pub async fn invoke_with_documents(
149        &self,
150        documents: Vec<Document>,
151        input: &str,
152    ) -> Result<String, ChainError>
153    where
154        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
155    {
156        if documents.is_empty() {
157            return Err(ChainError::ExecutionError(
158                "Document list is empty".to_string(),
159            ));
160        }
161
162        if self.verbose {
163            println!("\n=== MapReduceDocumentsChain ===");
164            println!("Document count: {}", documents.len());
165            println!("Input: {}", input);
166        }
167
168        if self.verbose {
169            println!("\n--- Map phase ---");
170        }
171
172        let mut map_futures = Vec::new();
173        for (i, doc) in documents.iter().enumerate() {
174            map_futures.push(self.map_document(doc, input, i));
175        }
176        let summaries: Vec<String> = try_join_all(map_futures).await?;
177
178        if self.verbose {
179            println!("\n--- Reduce phase ---");
180        }
181
182        let reduce_prompt = self.build_reduce_prompt(&summaries, input);
183
184        if self.verbose {
185            println!("Merging answers from {} documents", summaries.len());
186        }
187
188        let messages = vec![Message::human(&reduce_prompt)];
189        let response = self
190            .llm
191            .invoke(messages, None)
192            .await
193            .map_err(|e| ChainError::ExecutionError(format!("Reduce call failed: {}", e)))?;
194
195        let final_answer = response.content;
196
197        if self.verbose {
198            println!("Final answer: {}", final_answer);
199            println!("=== MapReduceDocumentsChain complete ===\n");
200        }
201
202        Ok(final_answer)
203    }
204}
205
206#[async_trait]
207impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for MapReduceDocumentsChain<M>
208where
209    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
210{
211    fn input_keys(&self) -> Vec<&str> {
212        vec![&self.input_key, "documents"]
213    }
214
215    fn output_keys(&self) -> Vec<&str> {
216        vec![&self.output_key]
217    }
218
219    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
220        let input = inputs
221            .get(&self.input_key)
222            .and_then(|v| v.as_str())
223            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
224
225        let documents: Vec<Document> = inputs
226            .get("documents")
227            .and_then(|v| v.as_array())
228            .map(|arr| {
229                arr.iter()
230                    .filter_map(|v| serde_json::from_value(v.clone()).ok())
231                    .collect()
232            })
233            .ok_or_else(|| ChainError::MissingInput("documents".to_string()))?;
234
235        let output = self.invoke_with_documents(documents, input).await?;
236
237        let mut result = HashMap::new();
238        result.insert(self.output_key.clone(), Value::String(output));
239        Ok(result)
240    }
241
242    fn name(&self) -> &str {
243        &self.name
244    }
245}
lc_chains/document_chains/map_reduce.rs

lc_chains/document_chains/
map_reduce.rs
