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

lc_chains/document_chains/
map_reduce.rs
