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 futures_util::StreamExt;
7use lc_core::language_models::LLMResult;
8use lc_core::{BaseChatModel, Runnable};
9use lc_schema::Message;
10use lc_shared::document::Document;
11use serde_json::Value;
12use std::collections::HashMap;
13
14use crate::base::{BaseChain, ChainError, ChainResult, ChainStream, StreamToken};
15
16/// Default Map processing prompt template.
17pub(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.
18
19Document content:
20{context}
21
22Question: {input}
23
24Answer based on this document:";
25
26/// Default Reduce merge prompt template.
27pub(crate) const DEFAULT_REDUCE_PROMPT: &str = "Below are answers from multiple documents. Please merge them into a single complete and coherent final answer.
28
29Answers from each document:
30{summaries}
31
32Original question: {input}
33
34Final consolidated answer:";
35
36/// MapReduceDocumentsChain
37///
38/// Processes documents in two steps:
39/// 1. Map: Calls LLM independently for each document to generate an answer
40/// 2. Reduce: Merges all independent answers into a final answer
41pub struct MapReduceDocumentsChain<M: BaseChatModel> {
42    llm: M,
43    map_prompt_template: String,
44    reduce_prompt_template: String,
45    document_variable_name: String,
46    input_key: String,
47    output_key: String,
48    name: String,
49    verbose: bool,
50}
51
52impl<M: BaseChatModel> MapReduceDocumentsChain<M> {
53    pub fn new(llm: M) -> Self {
54        Self {
55            llm,
56            map_prompt_template: DEFAULT_MAP_PROMPT.to_string(),
57            reduce_prompt_template: DEFAULT_REDUCE_PROMPT.to_string(),
58            document_variable_name: "context".to_string(),
59            input_key: "input".to_string(),
60            output_key: "output".to_string(),
61            name: "map_reduce_documents".to_string(),
62            verbose: false,
63        }
64    }
65
66    pub fn with_map_prompt(mut self, template: impl Into<String>) -> Self {
67        self.map_prompt_template = template.into();
68        self
69    }
70
71    pub fn with_reduce_prompt(mut self, template: impl Into<String>) -> Self {
72        self.reduce_prompt_template = template.into();
73        self
74    }
75
76    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
77        self.document_variable_name = name.into();
78        self
79    }
80
81    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
82        self.input_key = key.into();
83        self
84    }
85
86    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
87        self.output_key = key.into();
88        self
89    }
90
91    pub fn with_name(mut self, name: impl Into<String>) -> Self {
92        self.name = name.into();
93        self
94    }
95
96    pub fn with_verbose(mut self, verbose: bool) -> Self {
97        self.verbose = verbose;
98        self
99    }
100
101    pub fn build_map_prompt(&self, context: &str, input: &str) -> String {
102        self.map_prompt_template
103            .replace(&format!("{{{}}}", self.document_variable_name), context)
104            .replace("{input}", input)
105    }
106
107    pub fn build_reduce_prompt(&self, summaries: &[String], input: &str) -> String {
108        let summaries_text = summaries
109            .iter()
110            .enumerate()
111            .map(|(i, s)| format!("Answer from document {}:\n{}", i + 1, s))
112            .collect::<Vec<_>>()
113            .join("\n\n");
114
115        self.reduce_prompt_template
116            .replace("{summaries}", &summaries_text)
117            .replace("{input}", input)
118    }
119
120    /// Map phase: call LLM for a single document.
121    async fn map_document(
122        &self,
123        doc: &Document,
124        input: &str,
125        index: usize,
126    ) -> Result<String, ChainError>
127    where
128        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
129    {
130        let prompt = self.build_map_prompt(&doc.content, input);
131
132        if self.verbose {
133            println!("\n--- Map document {} ---", index + 1);
134        }
135
136        let messages = vec![Message::human(&prompt)];
137        let response = self.llm.invoke(messages, None).await.map_err(|e| {
138            ChainError::ExecutionError(format!("Map call failed (document {}): {}", index + 1, e))
139        })?;
140
141        if self.verbose {
142            println!("Document {} answer: {}", index + 1, response.content);
143        }
144
145        Ok(response.content)
146    }
147
148    /// Invoke with documents and input directly.
149    pub async fn invoke_with_documents(
150        &self,
151        documents: Vec<Document>,
152        input: &str,
153    ) -> Result<String, ChainError>
154    where
155        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
156    {
157        if documents.is_empty() {
158            return Err(ChainError::ExecutionError(
159                "Document list is empty".to_string(),
160            ));
161        }
162
163        if self.verbose {
164            println!("\n=== MapReduceDocumentsChain ===");
165            println!("Document count: {}", documents.len());
166            println!("Input: {}", input);
167        }
168
169        if self.verbose {
170            println!("\n--- Map phase ---");
171        }
172
173        let mut map_futures = Vec::new();
174        for (i, doc) in documents.iter().enumerate() {
175            map_futures.push(self.map_document(doc, input, i));
176        }
177        let summaries: Vec<String> = try_join_all(map_futures).await?;
178
179        if self.verbose {
180            println!("\n--- Reduce phase ---");
181        }
182
183        let reduce_prompt = self.build_reduce_prompt(&summaries, input);
184
185        if self.verbose {
186            println!("Merging answers from {} documents", summaries.len());
187        }
188
189        let messages = vec![Message::human(&reduce_prompt)];
190        let response = self
191            .llm
192            .invoke(messages, None)
193            .await
194            .map_err(|e| ChainError::ExecutionError(format!("Reduce call failed: {}", e)))?;
195
196        let final_answer = response.content;
197
198        if self.verbose {
199            println!("Final answer: {}", final_answer);
200            println!("=== MapReduceDocumentsChain complete ===\n");
201        }
202
203        Ok(final_answer)
204    }
205}
206
207#[async_trait]
208impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for MapReduceDocumentsChain<M>
209where
210    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
211{
212    fn input_keys(&self) -> Vec<&str> {
213        vec![&self.input_key, "documents"]
214    }
215
216    fn output_keys(&self) -> Vec<&str> {
217        vec![&self.output_key]
218    }
219
220    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
221        let input = inputs
222            .get(&self.input_key)
223            .and_then(|v| v.as_str())
224            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
225
226        let documents: Vec<Document> = inputs
227            .get("documents")
228            .and_then(|v| v.as_array())
229            .map(|arr| {
230                arr.iter()
231                    .filter_map(|v| serde_json::from_value(v.clone()).ok())
232                    .collect()
233            })
234            .ok_or_else(|| ChainError::MissingInput("documents".to_string()))?;
235
236        let output = self.invoke_with_documents(documents, input).await?;
237
238        let mut result = HashMap::new();
239        result.insert(self.output_key.clone(), Value::String(output));
240        Ok(result)
241    }
242
243    /// Stream execution for MapReduceDocumentsChain.
244    ///
245    /// The map phase runs via invoke (parallel, non-streaming, since reduce
246    /// needs all summaries). The reduce phase is streamed token by token
247    /// via `stream_chat`.
248    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
249        self.validate_inputs(&inputs)?;
250
251        let input = inputs
252            .get(&self.input_key)
253            .and_then(|v| v.as_str())
254            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
255
256        let documents: Vec<Document> = inputs
257            .get("documents")
258            .and_then(|v| v.as_array())
259            .map(|arr| {
260                arr.iter()
261                    .filter_map(|v| serde_json::from_value(v.clone()).ok())
262                    .collect()
263            })
264            .ok_or_else(|| ChainError::MissingInput("documents".to_string()))?;
265
266        if documents.is_empty() {
267            return Err(ChainError::ExecutionError(
268                "Document list is empty".to_string(),
269            ));
270        }
271
272        // Map phase: run all map calls in parallel (non-streaming)
273        let mut map_futures = Vec::new();
274        for (i, doc) in documents.iter().enumerate() {
275            map_futures.push(self.map_document(doc, input, i));
276        }
277        let summaries: Vec<String> = try_join_all(map_futures).await?;
278
279        // Reduce phase: stream the final merged answer
280        let reduce_prompt = self.build_reduce_prompt(&summaries, input);
281        let messages = vec![Message::human(&reduce_prompt)];
282
283        let llm_stream = self
284            .llm
285            .stream_chat(messages, None)
286            .await
287            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
288
289        let stream = llm_stream.map(|result| match result {
290            Ok(token) => Ok(StreamToken {
291                token,
292                is_final: false,
293            }),
294            Err(e) => Err(ChainError::StreamError(format!(
295                "Stream token error: {}",
296                e
297            ))),
298        });
299
300        let final_stream =
301            stream.chain(futures_util::stream::once(async move {
302                Ok(StreamToken {
303                    token: String::new(),
304                    is_final: true,
305                })
306            }));
307
308        Ok(Box::pin(final_stream))
309    }
310
311    fn name(&self) -> &str {
312        &self.name
313    }
314}
lc_chains/document_chains/map_reduce.rs

lc_chains/document_chains/
map_reduce.rs
