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::language_models::LLMResult;
8use lc_core::{BaseChatModel, Runnable};
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};
17
18/// Default Map + Rerank prompt template.
19pub(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).
20
21Document content:
22{context}
23
24Question: {input}
25
26Please output in the following format:
27Relevance score: <score>
28Answer: <your answer>";
29
30/// MapRerankDocumentsChain
31///
32/// First calls LLM independently for each document to generate an answer and score,
33/// then ranks by relevance score and returns the highest-scoring answer.
34pub struct MapRerankDocumentsChain<M: BaseChatModel> {
35    llm: M,
36    map_prompt_template: String,
37    document_variable_name: String,
38    input_key: String,
39    output_key: String,
40    name: String,
41    verbose: bool,
42    /// Return top k results (default 1, i.e. only the highest score).
43    top_k: usize,
44    /// Fallback score for LLM output without a parseable score (P1-3).
45    /// `None` = skip the document; `Some(n)` = rank it with score n.
46    default_score: Option<u32>,
47}
48
49// Pre-compiled regex patterns for score extraction.
50static SCORE_RE: LazyLock<Regex> =
51    LazyLock::new(|| Regex::new(r"(?i)(?:relevance\s*score|相关性评分)\s*[:：]\s*(\d+)").unwrap());
52static SCORE_RE2: LazyLock<Regex> =
53    LazyLock::new(|| Regex::new(r"(?i)score\s*[:：]\s*(\d+)").unwrap());
54
55/// Truncate a string to at most `max_len` characters, respecting char boundaries.
56fn truncate_str(s: &str, max_len: usize) -> &str {
57    if s.chars().count() <= max_len {
58        s
59    } else {
60        let end = s
61            .char_indices()
62            .nth(max_len)
63            .map(|(i, _)| i)
64            .unwrap_or(s.len());
65        &s[..end]
66    }
67}
68
69/// Extract score and answer from LLM output.
70///
71/// Returns `None` when no parseable score is present — the caller then decides
72/// how to treat unscored output (skip the document or use a configured default)
73/// instead of silently assigning a middle score that pollutes ranking (P1-3).
74pub fn extract_score(text: &str) -> Option<(u32, String)> {
75    for re in [&*SCORE_RE, &*SCORE_RE2] {
76        if let Some(caps) = re.captures(text) {
77            if let Ok(score) = caps[1].parse::<u32>() {
78                let cleaned = re.replace(text, "").trim().to_string();
79                let cleaned = cleaned
80                    .trim_start_matches("Answer")
81                    .trim_start_matches("答案")
82                    .trim_start_matches(&[':', '：'][..])
83                    .trim()
84                    .to_string();
85                return Some((
86                    std::cmp::min(score, 100),
87                    if cleaned.is_empty() {
88                        text.to_string()
89                    } else {
90                        cleaned
91                    },
92                ));
93            }
94        }
95    }
96    None
97}
98
99impl<M: BaseChatModel> MapRerankDocumentsChain<M> {
100    pub fn new(llm: M) -> Self {
101        Self {
102            llm,
103            map_prompt_template: DEFAULT_MAP_RERANK_PROMPT.to_string(),
104            document_variable_name: "context".to_string(),
105            input_key: "input".to_string(),
106            output_key: "output".to_string(),
107            name: "map_rerank_documents".to_string(),
108            verbose: false,
109            top_k: 1,
110            default_score: None,
111        }
112    }
113
114    pub fn with_map_prompt(mut self, template: impl Into<String>) -> Self {
115        self.map_prompt_template = template.into();
116        self
117    }
118
119    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
120        self.document_variable_name = name.into();
121        self
122    }
123
124    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
125        self.input_key = key.into();
126        self
127    }
128
129    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
130        self.output_key = key.into();
131        self
132    }
133
134    pub fn with_name(mut self, name: impl Into<String>) -> Self {
135        self.name = name.into();
136        self
137    }
138
139    pub fn with_verbose(mut self, verbose: bool) -> Self {
140        self.verbose = verbose;
141        self
142    }
143
144    /// Set the number of top results to return.
145    pub fn with_top_k(mut self, k: usize) -> Self {
146        self.top_k = k;
147        self
148    }
149
150    /// Configure the fallback score for LLM output without a parseable score.
151    ///
152    /// `None` (default) skips such documents; `Some(n)` ranks them with score `n`
153    /// instead of silently assigning the old middle score 50 (P1-3).
154    pub fn with_default_score(mut self, score: u32) -> Self {
155        self.default_score = Some(score);
156        self
157    }
158
159    /// Build Map stage prompt.
160    pub fn build_map_prompt(&self, context: &str, input: &str) -> String {
161        self.map_prompt_template
162            .replace(&format!("{{{}}}", self.document_variable_name), context)
163            .replace("{input}", input)
164    }
165
166    async fn map_document(
167        &self,
168        doc: &Document,
169        input: &str,
170        index: usize,
171    ) -> Result<Option<(u32, String)>, ChainError>
172    where
173        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
174    {
175        let prompt = self.build_map_prompt(&doc.content, input);
176        if self.verbose {
177            println!("\n--- Map document {} ---", index + 1);
178        }
179        let messages = vec![Message::human(&prompt)];
180        let response = self.llm.invoke(messages, None).await.map_err(|e| {
181            ChainError::ExecutionError(format!("Map call failed (document {}): {}", index + 1, e))
182        })?;
183
184        self.rank_output(&response.content, index)
185    }
186
187    /// Score an LLM output for one document: parse the relevance score and
188    /// answer, or apply the configured default / skip (P1-3). Shared by the
189    /// invoke and streaming map paths so the scoring rules never drift.
190    fn rank_output(&self, output: &str, index: usize) -> Result<Option<(u32, String)>, ChainError> {
191        // P1-3: no more silent middle score 50. Unscored output either uses the
192        // configured default_score or is excluded from ranking entirely.
193        let scored = match extract_score(output) {
194            Some(pair) => Some(pair),
195            None => match self.default_score {
196                Some(n) => Some((n, output.trim().to_string())),
197                None => {
198                    log::warn!(
199                        "MapRerank: document {} output has no parseable score; excluded from ranking",
200                        index + 1
201                    );
202                    None
203                }
204            },
205        };
206
207        if self.verbose {
208            if let Some((score, answer)) = &scored {
209                println!(
210                    "Document {} score: {}, answer: {}",
211                    index + 1,
212                    score,
213                    truncate_str(answer, 80)
214                );
215            } else {
216                println!("Document {} excluded (no score)", index + 1);
217            }
218        }
219        Ok(scored)
220    }
221
222    /// Map-phase variant for the streaming path: tokens are collected from
223    /// `stream_chat` so the full document answer is available for scoring
224    /// (ranking requires the complete output).
225    async fn map_document_stream(
226        &self,
227        doc: &Document,
228        input: &str,
229        index: usize,
230    ) -> Result<Option<(u32, String)>, ChainError>
231    where
232        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
233    {
234        let prompt = self.build_map_prompt(&doc.content, input);
235        if self.verbose {
236            println!("\n--- Map document {} (stream) ---", index + 1);
237        }
238        let messages = vec![Message::human(&prompt)];
239        let mut llm_stream = self.llm.stream_chat(messages, None).await.map_err(|e| {
240            ChainError::StreamError(format!("Map stream failed (document {}): {}", index + 1, e))
241        })?;
242
243        let mut text = String::new();
244        while let Some(chunk) = llm_stream.next().await {
245            match chunk {
246                Ok(token) => text.push_str(&token),
247                Err(e) => {
248                    return Err(ChainError::StreamError(format!(
249                        "Map stream token error (document {}): {}",
250                        index + 1,
251                        e
252                    )));
253                }
254            }
255        }
256        self.rank_output(&text, index)
257    }
258
259    /// Invoke with documents and input directly.
260    pub async fn invoke_with_documents(
261        &self,
262        documents: Vec<Document>,
263        input: &str,
264    ) -> Result<Vec<(u32, String)>, ChainError>
265    where
266        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
267    {
268        if documents.is_empty() {
269            return Err(ChainError::ExecutionError(
270                "Document list is empty".to_string(),
271            ));
272        }
273
274        if self.verbose {
275            println!("\n=== MapRerankDocumentsChain ===");
276            println!("Document count: {}, Input: {}", documents.len(), input);
277            println!("\n--- Map phase ---");
278        }
279
280        let mut map_futures = Vec::new();
281        for (i, doc) in documents.iter().enumerate() {
282            map_futures.push(self.map_document(doc, input, i));
283        }
284        // P1-3: drop documents whose output carried no score and had no default.
285        let mut results: Vec<(u32, String)> = try_join_all(map_futures)
286            .await?
287            .into_iter()
288            .flatten()
289            .collect();
290
291        if results.is_empty() {
292            return Err(ChainError::ExecutionError(
293                "All documents were excluded: no document produced a parseable score".to_string(),
294            ));
295        }
296
297        results.sort_by(|a, b| b.0.cmp(&a.0));
298
299        if self.verbose {
300            println!("\n--- Rerank phase ---");
301            for (i, (score, answer)) in results.iter().enumerate() {
302                println!(
303                    "Rank {}: score={}, answer={}",
304                    i + 1,
305                    score,
306                    truncate_str(answer, 100)
307                );
308            }
309        }
310
311        let top_results: Vec<(u32, String)> = results.into_iter().take(self.top_k).collect();
312        if self.verbose {
313            println!("Selected {} best results", top_results.len());
314            println!("=== MapRerankDocumentsChain complete ===\n");
315        }
316        Ok(top_results)
317    }
318}
319
320#[async_trait]
321impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for MapRerankDocumentsChain<M>
322where
323    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
324{
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;
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    #[derive(Debug)]
429    struct MockError(String);
430    impl std::fmt::Display for MockError {
431        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
432            write!(f, "{}", self.0)
433        }
434    }
435    impl std::error::Error for MockError {}
436
437    struct MockLLM;
438
439    #[async_trait]
440    impl Runnable<Vec<Message>, LLMResult> for MockLLM {
441        type Error = MockError;
442        async fn invoke(
443            &self,
444            _input: Vec<Message>,
445            _config: Option<RunnableConfig>,
446        ) -> Result<LLMResult, Self::Error> {
447            Ok(LLMResult {
448                content: "Relevance score: 90\nAnswer: best answer".to_string(),
449                model: "mock".to_string(),
450                token_usage: None,
451                tool_calls: None,
452                thinking_content: None,
453            })
454        }
455    }
456
457    #[async_trait]
458    impl BaseLanguageModel<Vec<Message>, LLMResult> for MockLLM {
459        fn model_name(&self) -> &str {
460            "mock"
461        }
462        fn get_num_tokens(&self, t: &str) -> usize {
463            t.len()
464        }
465        fn with_temperature(self, _: f32) -> Self {
466            self
467        }
468        fn with_max_tokens(self, _: usize) -> Self {
469            self
470        }
471    }
472
473    #[async_trait]
474    impl BaseChatModel for MockLLM {
475        async fn chat(
476            &self,
477            _messages: Vec<Message>,
478            _config: Option<RunnableConfig>,
479        ) -> Result<LLMResult, Self::Error> {
480            Ok(LLMResult {
481                content: "Relevance score: 90\nAnswer: best answer".to_string(),
482                model: "mock".to_string(),
483                token_usage: None,
484                tool_calls: None,
485                thinking_content: None,
486            })
487        }
488        async fn stream_chat(
489            &self,
490            _messages: Vec<Message>,
491            _config: Option<RunnableConfig>,
492        ) -> Result<Pin<Box<dyn Stream<Item = Result<String, Self::Error>> + Send>>, Self::Error>
493        {
494            let tokens = [Ok("Relevance score: 90\nAnswer: best answer".to_string())];
495            Ok(Box::pin(futures_util::stream::iter(tokens)))
496        }
497    }
498
499    #[tokio::test]
500    async fn test_map_rerank_stream_emits_top_answer() {
501        let chain = MapRerankDocumentsChain::new(MockLLM);
502        let docs = vec![Document::new("doc one"), Document::new("doc two")];
503        let docs_value = serde_json::to_value(docs).unwrap();
504        let mut inputs = HashMap::new();
505        inputs.insert("input".to_string(), Value::String("question".to_string()));
506        inputs.insert("documents".to_string(), docs_value);
507
508        let mut stream = chain.stream(inputs).await.unwrap();
509        let mut tokens = Vec::new();
510        while let Some(item) = stream.next().await {
511            tokens.push(item.unwrap());
512        }
513        assert_eq!(tokens.len(), 1);
514        assert!(tokens[0].is_final);
515        assert!(
516            tokens[0].token.contains("best answer"),
517            "top answer should be streamed, got {:?}",
518            tokens[0].token
519        );
520    }
521
522    #[tokio::test]
523    async fn test_map_rerank_stream_empty_documents() {
524        let chain = MapRerankDocumentsChain::new(MockLLM);
525        let mut inputs = HashMap::new();
526        inputs.insert("input".to_string(), Value::String("q".to_string()));
527        inputs.insert("documents".to_string(), serde_json::json!([]));
528        let err = match chain.stream(inputs).await {
529            Ok(_) => panic!("expected an execution error"),
530            Err(e) => e,
531        };
532        assert!(matches!(err, ChainError::ExecutionError(_)));
533    }
534
535    #[test]
536    fn test_rank_output_uses_default_score_for_unscored() {
537        let chain = MapRerankDocumentsChain::new(MockLLM).with_default_score(40);
538        let scored = chain.rank_output("plain answer without score", 0).unwrap();
539        assert_eq!(scored, Some((40, "plain answer without score".to_string())));
540    }
541
542    #[test]
543    fn test_rank_output_skips_unscored_without_default() {
544        let chain = MapRerankDocumentsChain::new(MockLLM);
545        assert_eq!(chain.rank_output("no score here", 0).unwrap(), None);
546    }
547}
lc_chains/document_chains/map_rerank.rs

lc_chains/document_chains/
map_rerank.rs
