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

lc_chains/document_chains/
map_rerank.rs
