langchain_rust/chain/combine_documents/
map_rerank.rs

1use std::{collections::HashMap, sync::LazyLock};
2
3use futures::{stream, StreamExt, TryStreamExt};
4use regex::Regex;
5use serde::{Deserialize, Serialize};
6use serde_json::{json, Value};
7
8use crate::{
9    chain::{
10        stuff_documents::{
11            input_documents, COMBINE_DOCUMENTS_DEFAULT_DOCUMENT_VARIABLE_NAME,
12            COMBINE_DOCUMENTS_DEFAULT_INPUT_KEY,
13        },
14        Chain, ChainError, DEFAULT_OUTPUT_KEY, DEFAULT_RESULT_KEY,
15    },
16    language_models::{GenerateResult, TokenUsage},
17    prompt::PromptArgs,
18};
19
20use super::refine::{add_usage, shared_inputs};
21
22const DEFAULT_MAX_CONCURRENCY: usize = 5;
23/// Output key with the winning answer's score (see [`MapRerankDocuments::execute`]).
24pub const MAP_RERANK_SCORE_KEY: &str = "score";
25/// Output key with every document's answer, when
26/// [`with_return_intermediate_steps`](MapRerankDocuments::with_return_intermediate_steps) is on.
27pub const MAP_RERANK_INTERMEDIATE_STEPS_KEY: &str = "intermediate_steps";
28
29/// The last `Score: N` in a model output. Tolerates markdown bold and decimals.
30static SCORE: LazyLock<Regex> = LazyLock::new(|| {
31    Regex::new(r"(?i)\**\s*score\s*\**\s*:\s*\**\s*(\d+(?:\.\d+)?)").expect("valid regex")
32});
33
34/// One document's answer, as parsed from the map chain's output.
35#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
36pub struct RankedAnswer {
37    /// The answer text (everything before `Score:`).
38    pub answer: String,
39    /// How well the answer covers the question, as the model rated it. `None` when the
40    /// output had no score.
41    pub score: Option<f64>,
42    /// Position of the source document in `input_documents`.
43    pub document_index: usize,
44}
45
46impl RankedAnswer {
47    /// Splits `Helpful Answer: ... \nScore: 80` style output. Reasoning in `<think>` tags
48    /// and anything before the last `Helpful Answer:` (models often echo the
49    /// `Question:` line) are dropped.
50    pub fn parse(output: &str, document_index: usize) -> Self {
51        let output = match output.rfind("</think>") {
52            Some(end) => &output[end + "</think>".len()..],
53            None => output,
54        };
55        let (answer, score) = match SCORE.captures_iter(output).last() {
56            Some(caps) => {
57                let start = caps.get(0).expect("whole match").start();
58                (&output[..start], caps[1].parse::<f64>().ok())
59            }
60            None => (output, None),
61        };
62        let answer = match answer.rfind("Helpful Answer:") {
63            Some(start) => &answer[start + "Helpful Answer:".len()..],
64            None => answer,
65        };
66        Self {
67            answer: answer.trim().to_string(),
68            score,
69            document_index,
70        }
71    }
72}
73
74/// Answers from every document separately, asking the model to score each answer, and
75/// returns the best one.
76///
77/// Good for questions whose answer sits in one document (a FAQ entry, a product page),
78/// where mixing documents would blur it. It makes one LLM call per document, concurrently.
79///
80/// Inputs: `input_documents` (a `Vec<Document>`) plus whatever the prompt uses (e.g.
81/// `question`). The chain gets each document's text as `context` and must end its output
82/// with `Score: <0-100>`, like the prompt of [`load_qa`](Self::load_qa) asks.
83///
84/// Outputs without a score rank below every scored one; if no output has a score the
85/// chain fails. Ties keep document order.
86///
87/// ```rust,ignore
88/// let chain = MapRerankDocuments::load_qa(OpenAI::default());
89/// let answer = chain
90///     .invoke(prompt_args! { "input_documents" => docs, "question" => "Who wrote it?" })
91///     .await?;
92/// ```
93pub struct MapRerankDocuments {
94    map_chain: Box<dyn Chain>,
95    input_key: String,
96    output_key: String,
97    document_variable_name: String,
98    max_concurrency: usize,
99    return_intermediate_steps: bool,
100}
101
102impl MapRerankDocuments {
103    /// `map_chain` runs once per document and must answer with a trailing `Score:`.
104    pub fn new<M: Into<Box<dyn Chain>>>(map_chain: M) -> Self {
105        Self {
106            map_chain: map_chain.into(),
107            input_key: COMBINE_DOCUMENTS_DEFAULT_INPUT_KEY.into(),
108            output_key: DEFAULT_OUTPUT_KEY.into(),
109            document_variable_name: COMBINE_DOCUMENTS_DEFAULT_DOCUMENT_VARIABLE_NAME.into(),
110            max_concurrency: DEFAULT_MAX_CONCURRENCY,
111            return_intermediate_steps: false,
112        }
113    }
114
115    /// Input key holding the documents (default `input_documents`).
116    pub fn with_input_key<S: Into<String>>(mut self, key: S) -> Self {
117        self.input_key = key.into();
118        self
119    }
120
121    /// Key of the answer in [`execute`](Chain::execute)'s output (default `output`).
122    pub fn with_output_key<S: Into<String>>(mut self, key: S) -> Self {
123        self.output_key = key.into();
124        self
125    }
126
127    /// Prompt variable each document's text goes in (default `context`).
128    pub fn with_document_variable_name<S: Into<String>>(mut self, name: S) -> Self {
129        self.document_variable_name = name.into();
130        self
131    }
132
133    /// How many LLM calls run at the same time (default 5, minimum 1).
134    pub fn with_max_concurrency(mut self, max_concurrency: usize) -> Self {
135        self.max_concurrency = max_concurrency.max(1);
136        self
137    }
138
139    /// Also return every document's [`RankedAnswer`] (in document order) from
140    /// [`execute`](Chain::execute), under [`MAP_RERANK_INTERMEDIATE_STEPS_KEY`].
141    pub fn with_return_intermediate_steps(mut self, enabled: bool) -> Self {
142        self.return_intermediate_steps = enabled;
143        self
144    }
145
146    /// Answers from every document and returns the answers in document order, the index
147    /// of the best one, and the usage of all calls.
148    pub async fn rank(
149        &self,
150        input_variables: &PromptArgs,
151    ) -> Result<(Vec<RankedAnswer>, usize, Option<TokenUsage>), ChainError> {
152        let docs = input_documents(input_variables, &self.input_key)?;
153        if docs.is_empty() {
154            return Err(ChainError::MissingInputVariable(format!(
155                "{} (no documents to rank)",
156                self.input_key
157            )));
158        }
159        let rest = shared_inputs(input_variables, &self.input_key);
160        let map_chain = self.map_chain.as_ref();
161        let map_inputs: Vec<PromptArgs> = docs
162            .iter()
163            .map(|doc| {
164                let mut inputs = rest.clone();
165                inputs.insert(
166                    self.document_variable_name.clone(),
167                    Value::String(doc.page_content.clone()),
168                );
169                inputs
170            })
171            .collect();
172        // `buffered` keeps the output in document order.
173        let results: Vec<GenerateResult> = stream::iter(map_inputs)
174            .map(|inputs| map_chain.call(inputs))
175            .buffered(self.max_concurrency)
176            .try_collect()
177            .await?;
178
179        let mut usage = None;
180        let answers: Vec<RankedAnswer> = results
181            .iter()
182            .enumerate()
183            .map(|(i, result)| {
184                add_usage(&mut usage, &result.tokens);
185                RankedAnswer::parse(&result.generation, i)
186            })
187            .collect();
188
189        let mut best: Option<(usize, f64)> = None;
190        for (i, answer) in answers.iter().enumerate() {
191            if let Some(score) = answer.score {
192                if best.is_none_or(|(_, top)| score > top) {
193                    best = Some((i, score));
194                }
195            }
196        }
197        let Some((best, _)) = best else {
198            let first = results[0].generation.chars().take(200).collect::<String>();
199            return Err(ChainError::OtherError(format!(
200                "map rerank: no answer had a `Score:` (first output: {first:?})"
201            )));
202        };
203        Ok((answers, best, usage))
204    }
205}
206
207#[async_trait::async_trait]
208impl Chain for MapRerankDocuments {
209    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
210        let (mut answers, best, tokens) = self.rank(&input_variables).await?;
211        Ok(GenerateResult {
212            generation: answers.swap_remove(best).answer,
213            tokens,
214            ..Default::default()
215        })
216    }
217
218    /// Returns the best answer under the output key, its score under
219    /// [`MAP_RERANK_SCORE_KEY`], and optionally every answer under
220    /// [`MAP_RERANK_INTERMEDIATE_STEPS_KEY`].
221    async fn execute(
222        &self,
223        input_variables: PromptArgs,
224    ) -> Result<HashMap<String, Value>, ChainError> {
225        let (answers, best, tokens) = self.rank(&input_variables).await?;
226        let winner = &answers[best];
227        let result = GenerateResult {
228            generation: winner.answer.clone(),
229            tokens,
230            ..Default::default()
231        };
232        let mut output = HashMap::new();
233        output.insert(self.output_key.clone(), json!(winner.answer));
234        output.insert(MAP_RERANK_SCORE_KEY.to_string(), json!(winner.score));
235        output.insert(DEFAULT_RESULT_KEY.to_string(), json!(result));
236        if self.return_intermediate_steps {
237            output.insert(
238                MAP_RERANK_INTERMEDIATE_STEPS_KEY.to_string(),
239                json!(answers),
240            );
241        }
242        Ok(output)
243    }
244
245    fn get_input_keys(&self) -> Vec<String> {
246        let mut keys = vec![self.input_key.clone()];
247        keys.extend(
248            self.map_chain
249                .get_input_keys()
250                .into_iter()
251                .filter(|k| *k != self.document_variable_name),
252        );
253        keys
254    }
255
256    fn get_output_keys(&self) -> Vec<String> {
257        let mut keys = vec![
258            self.output_key.clone(),
259            MAP_RERANK_SCORE_KEY.to_string(),
260            DEFAULT_RESULT_KEY.to_string(),
261        ];
262        if self.return_intermediate_steps {
263            keys.push(MAP_RERANK_INTERMEDIATE_STEPS_KEY.to_string());
264        }
265        keys
266    }
267}
langchain_rust/chain/combine_documents/map_rerank.rs

langchain_rust/chain/combine_documents/
map_rerank.rs
