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 lc_core::language_models::LLMResult;
7use lc_core::{BaseChatModel, Runnable};
8use lc_schema::Message;
9use lc_shared::document::Document;
10use regex::Regex;
11use serde_json::Value;
12use std::collections::HashMap;
13use std::sync::LazyLock;
14
15use crate::base::{BaseChain, ChainError, ChainResult};
16
17/// Default Map + Rerank prompt template.
18pub(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).
19
20Document content:
21{context}
22
23Question: {input}
24
25Please output in the following format:
26Relevance score: <score>
27Answer: <your answer>";
28
29/// MapRerankDocumentsChain
30///
31/// First calls LLM independently for each document to generate an answer and score,
32/// then ranks by relevance score and returns the highest-scoring answer.
33pub struct MapRerankDocumentsChain<M: BaseChatModel> {
34    llm: M,
35    map_prompt_template: String,
36    document_variable_name: String,
37    input_key: String,
38    output_key: String,
39    name: String,
40    verbose: bool,
41    /// Return top k results (default 1, i.e. only the highest score).
42    top_k: usize,
43}
44
45// Pre-compiled regex patterns for score extraction.
46static SCORE_RE: LazyLock<Regex> =
47    LazyLock::new(|| Regex::new(r"(?i)(?:relevance\s*score|相关性评分)\s*[:：]\s*(\d+)").unwrap());
48static SCORE_RE2: LazyLock<Regex> =
49    LazyLock::new(|| Regex::new(r"(?i)score\s*[:：]\s*(\d+)").unwrap());
50
51/// Truncate a string to at most `max_len` characters, respecting char boundaries.
52fn truncate_str(s: &str, max_len: usize) -> &str {
53    if s.chars().count() <= max_len {
54        s
55    } else {
56        let end = s
57            .char_indices()
58            .nth(max_len)
59            .map(|(i, _)| i)
60            .unwrap_or(s.len());
61        &s[..end]
62    }
63}
64
65/// Extract score and answer from LLM output.
66pub fn extract_score(text: &str) -> (u32, String) {
67    if let Some(caps) = SCORE_RE.captures(text) {
68        if let Ok(score) = caps[1].parse::<u32>() {
69            let cleaned = SCORE_RE.replace(text, "").trim().to_string();
70            let cleaned = cleaned
71                .trim_start_matches("Answer")
72                .trim_start_matches("答案")
73                .trim_start_matches(&[':', '：'][..])
74                .trim()
75                .to_string();
76            return (
77                std::cmp::min(score, 100),
78                if cleaned.is_empty() {
79                    text.to_string()
80                } else {
81                    cleaned
82                },
83            );
84        }
85    }
86
87    if let Some(caps) = SCORE_RE2.captures(text) {
88        if let Ok(score) = caps[1].parse::<u32>() {
89            let cleaned = SCORE_RE2.replace(text, "").trim().to_string();
90            return (
91                std::cmp::min(score, 100),
92                if cleaned.is_empty() {
93                    text.to_string()
94                } else {
95                    cleaned
96                },
97            );
98        }
99    }
100
101    (50, text.to_string())
102}
103
104impl<M: BaseChatModel> MapRerankDocumentsChain<M> {
105    pub fn new(llm: M) -> Self {
106        Self {
107            llm,
108            map_prompt_template: DEFAULT_MAP_RERANK_PROMPT.to_string(),
109            document_variable_name: "context".to_string(),
110            input_key: "input".to_string(),
111            output_key: "output".to_string(),
112            name: "map_rerank_documents".to_string(),
113            verbose: false,
114            top_k: 1,
115        }
116    }
117
118    pub fn with_map_prompt(mut self, template: impl Into<String>) -> Self {
119        self.map_prompt_template = template.into();
120        self
121    }
122
123    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
124        self.document_variable_name = name.into();
125        self
126    }
127
128    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
129        self.input_key = key.into();
130        self
131    }
132
133    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
134        self.output_key = key.into();
135        self
136    }
137
138    pub fn with_name(mut self, name: impl Into<String>) -> Self {
139        self.name = name.into();
140        self
141    }
142
143    pub fn with_verbose(mut self, verbose: bool) -> Self {
144        self.verbose = verbose;
145        self
146    }
147
148    /// Set the number of top results to return.
149    pub fn with_top_k(mut self, k: usize) -> Self {
150        self.top_k = k;
151        self
152    }
153
154    /// Build Map stage prompt.
155    pub fn build_map_prompt(&self, context: &str, input: &str) -> String {
156        self.map_prompt_template
157            .replace(&format!("{{{}}}", self.document_variable_name), context)
158            .replace("{input}", input)
159    }
160
161    async fn map_document(
162        &self,
163        doc: &Document,
164        input: &str,
165        index: usize,
166    ) -> Result<(u32, String), ChainError>
167    where
168        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
169    {
170        let prompt = self.build_map_prompt(&doc.content, input);
171        if self.verbose {
172            println!("\n--- Map document {} ---", index + 1);
173        }
174        let messages = vec![Message::human(&prompt)];
175        let response = self.llm.invoke(messages, None).await.map_err(|e| {
176            ChainError::ExecutionError(format!("Map call failed (document {}): {}", index + 1, e))
177        })?;
178        let (score, answer) = extract_score(&response.content);
179        if self.verbose {
180            println!(
181                "Document {} score: {}, answer: {}",
182                index + 1,
183                score,
184                truncate_str(&answer, 80)
185            );
186        }
187        Ok((score, answer))
188    }
189
190    /// Invoke with documents and input directly.
191    pub async fn invoke_with_documents(
192        &self,
193        documents: Vec<Document>,
194        input: &str,
195    ) -> Result<Vec<(u32, String)>, ChainError>
196    where
197        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
198    {
199        if documents.is_empty() {
200            return Err(ChainError::ExecutionError(
201                "Document list is empty".to_string(),
202            ));
203        }
204
205        if self.verbose {
206            println!("\n=== MapRerankDocumentsChain ===");
207            println!("Document count: {}, Input: {}", documents.len(), input);
208            println!("\n--- Map phase ---");
209        }
210
211        let mut map_futures = Vec::new();
212        for (i, doc) in documents.iter().enumerate() {
213            map_futures.push(self.map_document(doc, input, i));
214        }
215        let mut results: Vec<(u32, String)> = try_join_all(map_futures).await?;
216
217        results.sort_by(|a, b| b.0.cmp(&a.0));
218
219        if self.verbose {
220            println!("\n--- Rerank phase ---");
221            for (i, (score, answer)) in results.iter().enumerate() {
222                println!(
223                    "Rank {}: score={}, answer={}",
224                    i + 1,
225                    score,
226                    truncate_str(answer, 100)
227                );
228            }
229        }
230
231        let top_results: Vec<(u32, String)> = results.into_iter().take(self.top_k).collect();
232        if self.verbose {
233            println!("Selected {} best results", top_results.len());
234            println!("=== MapRerankDocumentsChain complete ===\n");
235        }
236        Ok(top_results)
237    }
238}
239
240#[async_trait]
241impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for MapRerankDocumentsChain<M>
242where
243    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
244{
245    fn input_keys(&self) -> Vec<&str> {
246        vec![&self.input_key, "documents"]
247    }
248    fn output_keys(&self) -> Vec<&str> {
249        vec![&self.output_key]
250    }
251
252    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
253        let input = inputs
254            .get(&self.input_key)
255            .and_then(|v| v.as_str())
256            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
257
258        let documents: Vec<Document> = inputs
259            .get("documents")
260            .and_then(|v| v.as_array())
261            .map(|arr| {
262                arr.iter()
263                    .filter_map(|v| serde_json::from_value(v.clone()).ok())
264                    .collect()
265            })
266            .ok_or_else(|| ChainError::MissingInput("documents".to_string()))?;
267
268        let results = self.invoke_with_documents(documents, input).await?;
269        let output_json: Vec<serde_json::Value> = results
270            .iter()
271            .map(|(score, answer)| serde_json::json!({"score": score, "answer": answer}))
272            .collect();
273
274        let mut result = HashMap::new();
275        result.insert(self.output_key.clone(), Value::Array(output_json));
276        Ok(result)
277    }
278
279    fn name(&self) -> &str {
280        &self.name
281    }
282}
lc_chains/document_chains/map_rerank.rs

lc_chains/document_chains/
map_rerank.rs
