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 futures_util::TryStreamExt;
8use lc_core::language_models::LLMResult;
9use lc_core::{BaseChatModel, Runnable};
10use lc_schema::Message;
11use lc_shared::document::Document;
12use serde_json::Value;
13use std::collections::HashMap;
14
15use crate::base::{BaseChain, ChainError, ChainResult, ChainStream, StreamToken};
16
17/// Default Map processing prompt template.
18pub(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.
19
20Document content:
21{context}
22
23Question: {input}
24
25Answer based on this document:";
26
27/// Default Reduce merge prompt template.
28pub(crate) const DEFAULT_REDUCE_PROMPT: &str = "Below are answers from multiple documents. Please merge them into a single complete and coherent final answer.
29
30Answers from each document:
31{summaries}
32
33Original question: {input}
34
35Final consolidated answer:";
36
37/// MapReduceDocumentsChain
38///
39/// Processes documents in two steps:
40/// 1. Map: Calls LLM independently for each document to generate an answer
41/// 2. Reduce: Merges all independent answers into a final answer
42pub struct MapReduceDocumentsChain<M: BaseChatModel> {
43    llm: M,
44    map_prompt_template: String,
45    reduce_prompt_template: String,
46    document_variable_name: String,
47    input_key: String,
48    output_key: String,
49    name: String,
50    verbose: bool,
51    /// Max number of in-flight map LLM calls; `None` = unbounded (P2-6).
52    map_concurrency: Option<usize>,
53}
54
55impl<M: BaseChatModel> MapReduceDocumentsChain<M> {
56    pub fn new(llm: M) -> Self {
57        Self {
58            llm,
59            map_prompt_template: DEFAULT_MAP_PROMPT.to_string(),
60            reduce_prompt_template: DEFAULT_REDUCE_PROMPT.to_string(),
61            document_variable_name: "context".to_string(),
62            input_key: "input".to_string(),
63            output_key: "output".to_string(),
64            name: "map_reduce_documents".to_string(),
65            verbose: false,
66            map_concurrency: None,
67        }
68    }
69
70    pub fn with_map_prompt(mut self, template: impl Into<String>) -> Self {
71        self.map_prompt_template = template.into();
72        self
73    }
74
75    pub fn with_reduce_prompt(mut self, template: impl Into<String>) -> Self {
76        self.reduce_prompt_template = template.into();
77        self
78    }
79
80    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
81        self.document_variable_name = name.into();
82        self
83    }
84
85    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
86        self.input_key = key.into();
87        self
88    }
89
90    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
91        self.output_key = key.into();
92        self
93    }
94
95    pub fn with_name(mut self, name: impl Into<String>) -> Self {
96        self.name = name.into();
97        self
98    }
99
100    pub fn with_verbose(mut self, verbose: bool) -> Self {
101        self.verbose = verbose;
102        self
103    }
104
105    /// Cap the number of concurrent map-phase LLM calls (P2-6).
106    ///
107    /// Default is unbounded (all documents mapped in parallel). Setting a
108    /// limit bounds in-flight requests — useful to respect provider rate
109    /// limits or bound memory on large document sets.
110    pub fn with_map_concurrency(mut self, limit: usize) -> Self {
111        self.map_concurrency = Some(limit);
112        self
113    }
114
115    pub fn build_map_prompt(&self, context: &str, input: &str) -> String {
116        self.map_prompt_template
117            .replace(&format!("{{{}}}", self.document_variable_name), context)
118            .replace("{input}", input)
119    }
120
121    pub fn build_reduce_prompt(&self, summaries: &[String], input: &str) -> String {
122        let summaries_text = summaries
123            .iter()
124            .enumerate()
125            .map(|(i, s)| format!("Answer from document {}:\n{}", i + 1, s))
126            .collect::<Vec<_>>()
127            .join("\n\n");
128
129        self.reduce_prompt_template
130            .replace("{summaries}", &summaries_text)
131            .replace("{input}", input)
132    }
133
134    /// Map phase: call LLM for a single document.
135    async fn map_document(
136        &self,
137        doc: &Document,
138        input: &str,
139        index: usize,
140    ) -> Result<String, ChainError>
141    where
142        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
143    {
144        let prompt = self.build_map_prompt(&doc.content, input);
145
146        if self.verbose {
147            println!("\n--- Map document {} ---", index + 1);
148        }
149
150        let messages = vec![Message::human(&prompt)];
151        let response = self.llm.invoke(messages, None).await.map_err(|e| {
152            ChainError::ExecutionError(format!("Map call failed (document {}): {}", index + 1, e))
153        })?;
154
155        if self.verbose {
156            println!("Document {} answer: {}", index + 1, response.content);
157        }
158
159        Ok(response.content)
160    }
161
162    /// Map phase: run the per-document LLM calls concurrently (P2-6).
163    ///
164    /// Default is unbounded parallelism (`try_join_all`, all documents in
165    /// flight at once). With [`Self::with_map_concurrency`] the number of
166    /// in-flight calls is capped via `buffer_unordered`, which bounds memory
167    /// and lets users respect provider rate limits.
168    async fn map_phase(
169        &self,
170        documents: &[Document],
171        input: &str,
172    ) -> Result<Vec<String>, ChainError>
173    where
174        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
175    {
176        // Materialize the per-document futures first: a `.map()` closure
177        // returning an async-fn future cannot express the higher-ranked
178        // `for<'a> FnMut((usize, &'a Document))` that `buffer_unordered`
179        // requires, so collect into a `Vec` of the concrete future type.
180        let mut map_futures = Vec::with_capacity(documents.len());
181        for (i, doc) in documents.iter().enumerate() {
182            map_futures.push(self.map_document(doc, input, i));
183        }
184
185        match self.map_concurrency {
186            Some(limit) => {
187                futures_util::stream::iter(map_futures)
188                    .buffer_unordered(limit)
189                    .try_collect()
190                    .await
191            }
192            None => try_join_all(map_futures).await,
193        }
194    }
195
196    /// Invoke with documents and input directly.
197    pub async fn invoke_with_documents(
198        &self,
199        documents: Vec<Document>,
200        input: &str,
201    ) -> Result<String, ChainError>
202    where
203        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
204    {
205        if documents.is_empty() {
206            return Err(ChainError::ExecutionError(
207                "Document list is empty".to_string(),
208            ));
209        }
210
211        if self.verbose {
212            println!("\n=== MapReduceDocumentsChain ===");
213            println!("Document count: {}", documents.len());
214            println!("Input: {}", input);
215        }
216
217        if self.verbose {
218            println!("\n--- Map phase ---");
219        }
220
221        let summaries = self.map_phase(&documents, input).await?;
222
223        if self.verbose {
224            println!("\n--- Reduce phase ---");
225        }
226
227        let reduce_prompt = self.build_reduce_prompt(&summaries, input);
228
229        if self.verbose {
230            println!("Merging answers from {} documents", summaries.len());
231        }
232
233        let messages = vec![Message::human(&reduce_prompt)];
234        let response = self
235            .llm
236            .invoke(messages, None)
237            .await
238            .map_err(|e| ChainError::ExecutionError(format!("Reduce call failed: {}", e)))?;
239
240        let final_answer = response.content;
241
242        if self.verbose {
243            println!("Final answer: {}", final_answer);
244            println!("=== MapReduceDocumentsChain complete ===\n");
245        }
246
247        Ok(final_answer)
248    }
249}
250
251#[async_trait]
252impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for MapReduceDocumentsChain<M>
253where
254    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
255{
256    fn input_keys(&self) -> Vec<&str> {
257        vec![&self.input_key, "documents"]
258    }
259
260    fn output_keys(&self) -> Vec<&str> {
261        vec![&self.output_key]
262    }
263
264    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
265        // P2-8: validate inputs on the invoke path too, matching stream.
266        self.validate_inputs(&inputs)?;
267
268        let input = inputs
269            .get(&self.input_key)
270            .and_then(|v| v.as_str())
271            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
272
273        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
274
275        let output = self.invoke_with_documents(documents, input).await?;
276
277        let mut result = HashMap::new();
278        result.insert(self.output_key.clone(), Value::String(output));
279        Ok(result)
280    }
281
282    /// Stream execution for MapReduceDocumentsChain.
283    ///
284    /// The map phase runs via invoke (parallel, non-streaming, since reduce
285    /// needs all summaries) — bounded by `map_concurrency` when configured.
286    /// The reduce phase is streamed token by token via `stream_chat`.
287    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
288        self.validate_inputs(&inputs)?;
289
290        let input = inputs
291            .get(&self.input_key)
292            .and_then(|v| v.as_str())
293            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
294
295        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
296
297        if documents.is_empty() {
298            return Err(ChainError::ExecutionError(
299                "Document list is empty".to_string(),
300            ));
301        }
302
303        // Map phase: run all map calls in parallel (non-streaming), bounded by
304        // `map_concurrency` when configured (P2-6).
305        let summaries = self.map_phase(&documents, input).await?;
306
307        // Reduce phase: stream the final merged answer
308        let reduce_prompt = self.build_reduce_prompt(&summaries, input);
309        let messages = vec![Message::human(&reduce_prompt)];
310
311        let llm_stream = self
312            .llm
313            .stream_chat(messages, None)
314            .await
315            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
316
317        let stream = llm_stream.map(|result| match result {
318            Ok(token) => Ok(StreamToken {
319                token,
320                is_final: false,
321            }),
322            Err(e) => Err(ChainError::StreamError(format!(
323                "Stream token error: {}",
324                e
325            ))),
326        });
327
328        let final_stream = stream.chain(futures_util::stream::once(async move {
329            Ok(StreamToken {
330                token: String::new(),
331                is_final: true,
332            })
333        }));
334
335        Ok(Box::pin(final_stream))
336    }
337
338    fn name(&self) -> &str {
339        &self.name
340    }
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346    use async_trait::async_trait;
347    use futures_util::Stream;
348    use lc_core::runnables::RunnableConfig;
349    use lc_core::{BaseLanguageModel, Runnable};
350    use std::pin::Pin;
351    use std::sync::atomic::{AtomicUsize, Ordering};
352    use std::sync::Arc;
353
354    #[derive(Debug)]
355    struct MockError(String);
356    impl std::fmt::Display for MockError {
357        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
358            write!(f, "{}", self.0)
359        }
360    }
361    impl std::error::Error for MockError {}
362
363    /// Tracking chat model that counts `invoke` calls and records the maximum
364    /// number of in-flight invokes, so the map phase's parallelism and the
365    /// `with_map_concurrency` cap are provable rather than assumed. Each invoke
366    /// yields repeatedly so concurrently-polled futures genuinely overlap.
367    struct TrackingLLM {
368        invokes: Arc<AtomicUsize>,
369        in_flight: Arc<AtomicUsize>,
370        max_in_flight: Arc<AtomicUsize>,
371    }
372
373    impl TrackingLLM {
374        fn counters() -> (Arc<AtomicUsize>, Arc<AtomicUsize>, Arc<AtomicUsize>) {
375            (
376                Arc::new(AtomicUsize::new(0)),
377                Arc::new(AtomicUsize::new(0)),
378                Arc::new(AtomicUsize::new(0)),
379            )
380        }
381    }
382
383    #[async_trait]
384    impl Runnable<Vec<Message>, LLMResult> for TrackingLLM {
385        type Error = MockError;
386        async fn invoke(
387            &self,
388            input: Vec<Message>,
389            _config: Option<RunnableConfig>,
390        ) -> Result<LLMResult, Self::Error> {
391            let cur = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
392            self.max_in_flight.fetch_max(cur, Ordering::SeqCst);
393
394            // Yield repeatedly so every concurrently-polled future is actually
395            // scheduled before any of them returns.
396            for _ in 0..32 {
397                tokio::task::yield_now().await;
398            }
399
400            self.in_flight.fetch_sub(1, Ordering::SeqCst);
401            self.invokes.fetch_add(1, Ordering::SeqCst);
402
403            let is_reduce = input
404                .iter()
405                .any(|m| m.content.contains("Below are answers"));
406            let content = if is_reduce {
407                "final merged answer".to_string()
408            } else {
409                "map answer".to_string()
410            };
411            Ok(LLMResult {
412                content,
413                model: "mock".to_string(),
414                token_usage: None,
415                tool_calls: None,
416                thinking_content: None,
417            })
418        }
419    }
420
421    #[async_trait]
422    impl BaseLanguageModel<Vec<Message>, LLMResult> for TrackingLLM {
423        fn model_name(&self) -> &str {
424            "mock"
425        }
426        fn get_num_tokens(&self, t: &str) -> usize {
427            t.len()
428        }
429        fn with_temperature(self, _: f32) -> Self {
430            self
431        }
432        fn with_max_tokens(self, _: usize) -> Self {
433            self
434        }
435    }
436
437    #[async_trait]
438    impl BaseChatModel for TrackingLLM {
439        async fn chat(
440            &self,
441            messages: Vec<Message>,
442            _config: Option<RunnableConfig>,
443        ) -> Result<LLMResult, Self::Error> {
444            // Same counting behavior as invoke so the stream path's map phase
445            // is observably identical.
446            self.invoke(messages, None).await
447        }
448        async fn stream_chat(
449            &self,
450            _messages: Vec<Message>,
451            _config: Option<RunnableConfig>,
452        ) -> Result<Pin<Box<dyn Stream<Item = Result<String, Self::Error>> + Send>>, Self::Error>
453        {
454            let tokens = [Ok("streamed token".to_string())];
455            Ok(Box::pin(futures_util::stream::iter(tokens)))
456        }
457    }
458
459    fn docs(n: usize) -> Vec<Document> {
460        (0..n).map(|i| Document::new(format!("doc {i}"))).collect()
461    }
462
463    /// P2-6: the map phase is genuinely parallel — with 6 documents the map
464    /// LLM calls overlap in flight (max in-flight > 1), then a single reduce
465    /// call produces the final answer.
466    #[tokio::test]
467    async fn test_map_reduce_map_phase_is_parallel() {
468        let (invokes, in_flight, max_in_flight) = TrackingLLM::counters();
469        let chain = MapReduceDocumentsChain::new(TrackingLLM {
470            invokes: invokes.clone(),
471            in_flight,
472            max_in_flight: max_in_flight.clone(),
473        });
474
475        let out = chain
476            .invoke_with_documents(docs(6), "question")
477            .await
478            .unwrap();
479        assert_eq!(out, "final merged answer");
480        assert_eq!(invokes.load(Ordering::SeqCst), 7, "6 map + 1 reduce");
481        assert!(
482            max_in_flight.load(Ordering::SeqCst) > 1,
483            "map phase must overlap calls, got max in-flight {}",
484            max_in_flight.load(Ordering::SeqCst)
485        );
486    }
487
488    /// P2-6: `with_map_concurrency(2)` caps in-flight map calls at 2 while
489    /// still running concurrently (max in-flight exactly 2).
490    #[tokio::test]
491    async fn test_map_reduce_concurrency_limit_caps_in_flight() {
492        let (invokes, in_flight, max_in_flight) = TrackingLLM::counters();
493        let chain = MapReduceDocumentsChain::new(TrackingLLM {
494            invokes: invokes.clone(),
495            in_flight,
496            max_in_flight: max_in_flight.clone(),
497        })
498        .with_map_concurrency(2);
499
500        let out = chain
501            .invoke_with_documents(docs(6), "question")
502            .await
503            .unwrap();
504        assert_eq!(out, "final merged answer");
505        assert_eq!(invokes.load(Ordering::SeqCst), 7);
506        assert_eq!(
507            max_in_flight.load(Ordering::SeqCst),
508            2,
509            "concurrency cap 2 must bound in-flight map calls"
510        );
511    }
512
513    /// P2-6: empty documents error loudly before any LLM call.
514    #[tokio::test]
515    async fn test_map_reduce_empty_documents() {
516        let (invokes, in_flight, max_in_flight) = TrackingLLM::counters();
517        let chain = MapReduceDocumentsChain::new(TrackingLLM {
518            invokes: invokes.clone(),
519            in_flight,
520            max_in_flight,
521        });
522        let err = match chain.invoke_with_documents(vec![], "q").await {
523            Ok(_) => panic!("expected an execution error"),
524            Err(e) => e,
525        };
526        assert!(matches!(err, ChainError::ExecutionError(_)));
527        assert_eq!(invokes.load(Ordering::SeqCst), 0);
528    }
529}
lc_chains/document_chains/map_reduce.rs

lc_chains/document_chains/
map_reduce.rs
