lc_chains/document_chains/
map_rerank.rs

1// lc-chains/src/document_chains/map_rerank.rs
2//! MapRerankDocumentsChain - processes documents in parallel then ranks by relevance.
3
4use async_trait::async_trait;
5use futures_util::future::try_join_all;
6use futures_util::StreamExt;
7use lc_core::BaseChatModel;
8use lc_providers::{wrap_chat_model, ProviderError};
9use lc_schema::Message;
10use lc_shared::document::Document;
11use regex::Regex;
12use serde_json::Value;
13use std::collections::HashMap;
14use std::sync::LazyLock;
15
16use crate::base::{BaseChain, ChainError, ChainResult, ChainStream, StreamToken};
17use crate::BoxedChatModel;
18
19/// Default Map + Rerank prompt template.
20pub(crate) const DEFAULT_MAP_RERANK_PROMPT: &str = "Answer the question based on the following document, and provide a relevance score (0-100, higher is more relevant).
21
22Document content:
23{context}
24
25Question: {input}
26
27Please output in the following format:
28Relevance score: <score>
29Answer: <your answer>";
30
31/// MapRerankDocumentsChain
32///
33/// First calls LLM independently for each document to generate an answer and score,
34/// then ranks by relevance score and returns the highest-scoring answer.
35pub struct MapRerankDocumentsChain {
36    llm: BoxedChatModel,
37    map_prompt_template: String,
38    document_variable_name: String,
39    input_key: String,
40    output_key: String,
41    name: String,
42    verbose: bool,
43    /// Return top k results (default 1, i.e. only the highest score).
44    top_k: usize,
45    /// Fallback score for LLM output without a parseable score (P1-3).
46    /// `None` = skip the document; `Some(n)` = rank it with score n.
47    default_score: Option<u32>,
48}
49
50// Pre-compiled regex patterns for score extraction.
51static SCORE_RE: LazyLock<Regex> =
52    LazyLock::new(|| Regex::new(r"(?i)(?:relevance\s*score|相关性评分)\s*[:：]\s*(\d+)").unwrap());
53static SCORE_RE2: LazyLock<Regex> =
54    LazyLock::new(|| Regex::new(r"(?i)score\s*[:：]\s*(\d+)").unwrap());
55
56/// Truncate a string to at most `max_len` characters, respecting char boundaries.
57fn truncate_str(s: &str, max_len: usize) -> &str {
58    if s.chars().count() <= max_len {
59        s
60    } else {
61        let end = s
62            .char_indices()
63            .nth(max_len)
64            .map(|(i, _)| i)
65            .unwrap_or(s.len());
66        &s[..end]
67    }
68}
69
70/// Extract score and answer from LLM output.
71///
72/// Returns `None` when no parseable score is present — the caller then decides
73/// how to treat unscored output (skip the document or use a configured default)
74/// instead of silently assigning a middle score that pollutes ranking (P1-3).
75pub fn extract_score(text: &str) -> Option<(u32, String)> {
76    for re in [&*SCORE_RE, &*SCORE_RE2] {
77        if let Some(caps) = re.captures(text) {
78            if let Ok(score) = caps[1].parse::<u32>() {
79                let cleaned = re.replace(text, "").trim().to_string();
80                let cleaned = cleaned
81                    .trim_start_matches("Answer")
82                    .trim_start_matches("答案")
83                    .trim_start_matches(&[':', '：'][..])
84                    .trim()
85                    .to_string();
86                return Some((
87                    std::cmp::min(score, 100),
88                    if cleaned.is_empty() {
89                        text.to_string()
90                    } else {
91                        cleaned
92                    },
93                ));
94            }
95        }
96    }
97    None
98}
99
100impl MapRerankDocumentsChain {
101    /// Create a new [`MapRerankDocumentsChain`] with the given LLM.
102    pub fn new<L>(llm: L) -> Self
103    where
104        L: BaseChatModel + Send + Sync + 'static,
105        L::Error: Into<ProviderError>,
106    {
107        Self {
108            llm: wrap_chat_model(llm),
109            map_prompt_template: DEFAULT_MAP_RERANK_PROMPT.to_string(),
110            document_variable_name: "context".to_string(),
111            input_key: "input".to_string(),
112            output_key: "output".to_string(),
113            name: "map_rerank_documents".to_string(),
114            verbose: false,
115            top_k: 1,
116            default_score: None,
117        }
118    }
119
120    /// Set the map-phase prompt template.
121    pub fn with_map_prompt(mut self, template: impl Into<String>) -> Self {
122        self.map_prompt_template = template.into();
123        self
124    }
125
126    /// Set the document variable name used in the map prompt.
127    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
128        self.document_variable_name = name.into();
129        self
130    }
131
132    /// Set the input key.
133    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
134        self.input_key = key.into();
135        self
136    }
137
138    /// Set the output key.
139    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
140        self.output_key = key.into();
141        self
142    }
143
144    /// Set the chain name.
145    pub fn with_name(mut self, name: impl Into<String>) -> Self {
146        self.name = name.into();
147        self
148    }
149
150    /// Set verbose mode.
151    pub fn with_verbose(mut self, verbose: bool) -> Self {
152        self.verbose = verbose;
153        self
154    }
155
156    /// Set the number of top results to return.
157    pub fn with_top_k(mut self, k: usize) -> Self {
158        self.top_k = k;
159        self
160    }
161
162    /// Configure the fallback score for LLM output without a parseable score.
163    ///
164    /// `None` (default) skips such documents; `Some(n)` ranks them with score `n`
165    /// instead of silently assigning the old middle score 50 (P1-3).
166    pub fn with_default_score(mut self, score: u32) -> Self {
167        self.default_score = Some(score);
168        self
169    }
170
171    /// Build Map stage prompt.
172    pub fn build_map_prompt(&self, context: &str, input: &str) -> String {
173        self.map_prompt_template
174            .replace(&format!("{{{}}}", self.document_variable_name), context)
175            .replace("{input}", input)
176    }
177
178    async fn map_document(
179        &self,
180        doc: &Document,
181        input: &str,
182        index: usize,
183    ) -> Result<Option<(u32, String)>, ChainError> {
184        let prompt = self.build_map_prompt(&doc.content, input);
185        if self.verbose {
186            println!("\n--- Map document {} ---", index + 1);
187        }
188        let messages = vec![Message::human(&prompt)];
189        let response = self.llm.invoke(messages, None).await.map_err(|e| {
190            ChainError::ExecutionError(format!("Map call failed (document {}): {}", index + 1, e))
191        })?;
192
193        self.rank_output(&response.content, index)
194    }
195
196    /// Score an LLM output for one document: parse the relevance score and
197    /// answer, or apply the configured default / skip (P1-3). Shared by the
198    /// invoke and streaming map paths so the scoring rules never drift.
199    fn rank_output(&self, output: &str, index: usize) -> Result<Option<(u32, String)>, ChainError> {
200        // P1-3: no more silent middle score 50. Unscored output either uses the
201        // configured default_score or is excluded from ranking entirely.
202        let scored = match extract_score(output) {
203            Some(pair) => Some(pair),
204            None => match self.default_score {
205                Some(n) => Some((n, output.trim().to_string())),
206                None => {
207                    log::warn!(
208                        "MapRerank: document {} output has no parseable score; excluded from ranking",
209                        index + 1
210                    );
211                    None
212                }
213            },
214        };
215
216        if self.verbose {
217            if let Some((score, answer)) = &scored {
218                println!(
219                    "Document {} score: {}, answer: {}",
220                    index + 1,
221                    score,
222                    truncate_str(answer, 80)
223                );
224            } else {
225                println!("Document {} excluded (no score)", index + 1);
226            }
227        }
228        Ok(scored)
229    }
230
231    /// Map-phase variant for the streaming path: tokens are collected from
232    /// `stream_chat` so the full document answer is available for scoring
233    /// (ranking requires the complete output).
234    async fn map_document_stream(
235        &self,
236        doc: &Document,
237        input: &str,
238        index: usize,
239    ) -> Result<Option<(u32, String)>, ChainError> {
240        let prompt = self.build_map_prompt(&doc.content, input);
241        if self.verbose {
242            println!("\n--- Map document {} (stream) ---", index + 1);
243        }
244        let messages = vec![Message::human(&prompt)];
245        let mut llm_stream = self.llm.stream_chat(messages, None).await.map_err(|e| {
246            ChainError::StreamError(format!("Map stream failed (document {}): {}", index + 1, e))
247        })?;
248
249        let mut text = String::new();
250        while let Some(chunk) = llm_stream.next().await {
251            match chunk {
252                Ok(chunk) => text.push_str(&chunk.text),
253                Err(e) => {
254                    return Err(ChainError::StreamError(format!(
255                        "Map stream token error (document {}): {}",
256                        index + 1,
257                        e
258                    )));
259                }
260            }
261        }
262        self.rank_output(&text, index)
263    }
264
265    /// Invoke with documents and input directly.
266    pub async fn invoke_with_documents(
267        &self,
268        documents: Vec<Document>,
269        input: &str,
270    ) -> Result<Vec<(u32, String)>, ChainError> {
271        if documents.is_empty() {
272            return Err(ChainError::ExecutionError(
273                "Document list is empty".to_string(),
274            ));
275        }
276
277        if self.verbose {
278            println!("\n=== MapRerankDocumentsChain ===");
279            println!("Document count: {}, Input: {}", documents.len(), input);
280            println!("\n--- Map phase ---");
281        }
282
283        let mut map_futures = Vec::new();
284        for (i, doc) in documents.iter().enumerate() {
285            map_futures.push(self.map_document(doc, input, i));
286        }
287        // P1-3: drop documents whose output carried no score and had no default.
288        let mut results: Vec<(u32, String)> = try_join_all(map_futures)
289            .await?
290            .into_iter()
291            .flatten()
292            .collect();
293
294        if results.is_empty() {
295            return Err(ChainError::ExecutionError(
296                "All documents were excluded: no document produced a parseable score".to_string(),
297            ));
298        }
299
300        results.sort_by(|a, b| b.0.cmp(&a.0));
301
302        if self.verbose {
303            println!("\n--- Rerank phase ---");
304            for (i, (score, answer)) in results.iter().enumerate() {
305                println!(
306                    "Rank {}: score={}, answer={}",
307                    i + 1,
308                    score,
309                    truncate_str(answer, 100)
310                );
311            }
312        }
313
314        let top_results: Vec<(u32, String)> = results.into_iter().take(self.top_k).collect();
315        if self.verbose {
316            println!("Selected {} best results", top_results.len());
317            println!("=== MapRerankDocumentsChain complete ===\n");
318        }
319        Ok(top_results)
320    }
321}
322
323#[async_trait]
324impl BaseChain for MapRerankDocumentsChain {
325    fn input_keys(&self) -> Vec<&str> {
326        vec![&self.input_key, "documents"]
327    }
328    fn output_keys(&self) -> Vec<&str> {
329        vec![&self.output_key]
330    }
331
332    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
333        // P2-8: validate inputs on the invoke path too, matching stream.
334        self.validate_inputs(&inputs)?;
335
336        let input = inputs
337            .get(&self.input_key)
338            .and_then(|v| v.as_str())
339            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
340
341        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
342
343        let results = self.invoke_with_documents(documents, input).await?;
344        let output_json: Vec<serde_json::Value> = results
345            .iter()
346            .map(|(score, answer)| serde_json::json!({"score": score, "answer": answer}))
347            .collect();
348
349        let mut result = HashMap::new();
350        result.insert(self.output_key.clone(), Value::Array(output_json));
351        Ok(result)
352    }
353
354    /// Stream execution for MapRerankDocumentsChain.
355    ///
356    /// P2-2: the map phase runs via `stream_chat` per document (tokens
357    /// accumulated for scoring), then the reranked top answer(s) are emitted.
358    /// Raw token streaming of the final answer is impossible here — ranking
359    /// requires each document's complete output — so the ranked result is the
360    /// stream payload, produced without the base default's silent `unwrap_or("")`.
361    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
362        // P2-8: validate inputs on the stream path too (this was the one
363        // document chain that skipped it entirely), matching the others.
364        self.validate_inputs(&inputs)?;
365
366        let input = inputs
367            .get(&self.input_key)
368            .and_then(|v| v.as_str())
369            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
370
371        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
372        if documents.is_empty() {
373            return Err(ChainError::ExecutionError(
374                "Document list is empty".to_string(),
375            ));
376        }
377
378        let mut map_futures = Vec::new();
379        for (i, doc) in documents.iter().enumerate() {
380            map_futures.push(self.map_document_stream(doc, input, i));
381        }
382        let mut results: Vec<(u32, String)> = try_join_all(map_futures)
383            .await?
384            .into_iter()
385            .flatten()
386            .collect();
387
388        if results.is_empty() {
389            return Err(ChainError::ExecutionError(
390                "All documents were excluded: no document produced a parseable score".to_string(),
391            ));
392        }
393
394        results.sort_by(|a, b| b.0.cmp(&a.0));
395        let top_results: Vec<(u32, String)> = results.into_iter().take(self.top_k).collect();
396
397        let stream = futures_util::stream::once(async move {
398            let text = top_results
399                .iter()
400                .map(|(_, answer)| answer.as_str())
401                .collect::<Vec<_>>()
402                .join("\n\n");
403            Ok(StreamToken {
404                token: text,
405                is_final: true,
406            })
407        });
408
409        Ok(Box::pin(stream))
410    }
411
412    fn name(&self) -> &str {
413        &self.name
414    }
415}
416
417#[cfg(test)]
418mod tests {
419    use super::*;
420    use async_trait::async_trait;
421    use futures_util::Stream;
422    use lc_core::language_models::{LLMResult, StreamChunk};
423    use lc_core::runnables::RunnableConfig;
424    use lc_core::{BaseLanguageModel, Runnable};
425    use std::pin::Pin;
426
427    /// Mock chat model whose stream returns a fixed scored answer per document.
428    struct MockLLM;
429
430    #[async_trait]
431    impl Runnable<Vec<Message>, LLMResult> for MockLLM {
432        type Error = ProviderError;
433        async fn invoke(
434            &self,
435            _input: Vec<Message>,
436            _config: Option<RunnableConfig>,
437        ) -> Result<LLMResult, Self::Error> {
438            Ok(LLMResult {
439                content: "Relevance score: 90\nAnswer: best answer".to_string(),
440                model: "mock".to_string(),
441                token_usage: None,
442                tool_calls: None,
443                thinking_content: None,
444            })
445        }
446    }
447
448    #[async_trait]
449    impl BaseLanguageModel<Vec<Message>, LLMResult> for MockLLM {
450        fn model_name(&self) -> &str {
451            "mock"
452        }
453        fn get_num_tokens(&self, t: &str) -> usize {
454            t.len()
455        }
456        fn with_temperature(self, _: f32) -> Self {
457            self
458        }
459        fn with_max_tokens(self, _: usize) -> Self {
460            self
461        }
462    }
463
464    #[async_trait]
465    impl BaseChatModel for MockLLM {
466        async fn chat(
467            &self,
468            _messages: Vec<Message>,
469            _config: Option<RunnableConfig>,
470        ) -> Result<LLMResult, Self::Error> {
471            Ok(LLMResult {
472                content: "Relevance score: 90\nAnswer: best answer".to_string(),
473                model: "mock".to_string(),
474                token_usage: None,
475                tool_calls: None,
476                thinking_content: None,
477            })
478        }
479        async fn stream_chat(
480            &self,
481            _messages: Vec<Message>,
482            _config: Option<RunnableConfig>,
483        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
484        {
485            let tokens = [Ok(StreamChunk::new(
486                "Relevance score: 90\nAnswer: best answer",
487            ))];
488            Ok(Box::pin(futures_util::stream::iter(tokens)))
489        }
490    }
491
492    #[tokio::test]
493    async fn test_map_rerank_stream_emits_top_answer() {
494        let chain = MapRerankDocumentsChain::new(MockLLM);
495        let docs = vec![Document::new("doc one"), Document::new("doc two")];
496        let docs_value = serde_json::to_value(docs).unwrap();
497        let mut inputs = HashMap::new();
498        inputs.insert("input".to_string(), Value::String("question".to_string()));
499        inputs.insert("documents".to_string(), docs_value);
500
501        let mut stream = chain.stream(inputs).await.unwrap();
502        let mut tokens = Vec::new();
503        while let Some(item) = stream.next().await {
504            tokens.push(item.unwrap());
505        }
506        assert_eq!(tokens.len(), 1);
507        assert!(tokens[0].is_final);
508        assert!(
509            tokens[0].token.contains("best answer"),
510            "top answer should be streamed, got {:?}",
511            tokens[0].token
512        );
513    }
514
515    #[tokio::test]
516    async fn test_map_rerank_stream_empty_documents() {
517        let chain = MapRerankDocumentsChain::new(MockLLM);
518        let mut inputs = HashMap::new();
519        inputs.insert("input".to_string(), Value::String("q".to_string()));
520        inputs.insert("documents".to_string(), serde_json::json!([]));
521        let err = match chain.stream(inputs).await {
522            Ok(_) => panic!("expected an execution error"),
523            Err(e) => e,
524        };
525        assert!(matches!(err, ChainError::ExecutionError(_)));
526    }
527
528    #[test]
529    fn test_rank_output_uses_default_score_for_unscored() {
530        let chain = MapRerankDocumentsChain::new(MockLLM).with_default_score(40);
531        let scored = chain.rank_output("plain answer without score", 0).unwrap();
532        assert_eq!(scored, Some((40, "plain answer without score".to_string())));
533    }
534
535    #[test]
536    fn test_rank_output_skips_unscored_without_default() {
537        let chain = MapRerankDocumentsChain::new(MockLLM);
538        assert_eq!(chain.rank_output("no score here", 0).unwrap(), None);
539    }
540}
lc_chains/document_chains/map_rerank.rs

lc_chains/document_chains/
map_rerank.rs
