lc_chains/document_chains/
refine.rs

1// lc-chains/src/document_chains/refine.rs
2//! RefineDocumentsChain - iteratively refines the answer document by document.
3
4use async_trait::async_trait;
5use futures_util::StreamExt;
6use lc_core::language_models::LLMResult;
7use lc_core::{BaseChatModel, Runnable};
8use lc_schema::Message;
9use lc_shared::document::Document;
10use serde_json::Value;
11use std::collections::HashMap;
12
13use crate::base::{BaseChain, ChainError, ChainResult, ChainStream, StreamToken};
14
15/// Default initial processing prompt template.
16pub(crate) const DEFAULT_REFINE_INITIAL_PROMPT: &str =
17    "Answer the question based on the following reference information.
18
19Reference information:
20{context}
21
22Question: {input}
23
24Answer:";
25
26/// Default iterative refinement prompt template.
27pub(crate) const DEFAULT_REFINE_PROMPT: &str = "You have provided an answer based on partial information. Here is additional reference information.
28
29Existing answer:
30{existing_answer}
31
32New reference information:
33{context}
34
35Please refine or modify your answer based on the new information. If the new information does not conflict with the existing answer, merge them. If the new information conflicts with the existing answer, prioritize the new information.
36
37Question: {input}
38
39Refined answer:";
40
41/// RefineDocumentsChain
42///
43/// Iteratively refines the answer document by document.
44/// Generates an initial answer from the first document, then refines with each subsequent document.
45pub struct RefineDocumentsChain<M: BaseChatModel> {
46    llm: M,
47    initial_prompt_template: String,
48    refine_prompt_template: String,
49    document_variable_name: String,
50    input_key: String,
51    output_key: String,
52    name: String,
53    verbose: bool,
54}
55
56impl<M: BaseChatModel> RefineDocumentsChain<M> {
57    pub fn new(llm: M) -> Self {
58        Self {
59            llm,
60            initial_prompt_template: DEFAULT_REFINE_INITIAL_PROMPT.to_string(),
61            refine_prompt_template: DEFAULT_REFINE_PROMPT.to_string(),
62            document_variable_name: "context".to_string(),
63            input_key: "input".to_string(),
64            output_key: "output".to_string(),
65            name: "refine_documents".to_string(),
66            verbose: false,
67        }
68    }
69
70    pub fn with_initial_prompt(mut self, template: impl Into<String>) -> Self {
71        self.initial_prompt_template = template.into();
72        self
73    }
74
75    pub fn with_refine_prompt(mut self, template: impl Into<String>) -> Self {
76        self.refine_prompt_template = template.into();
77        self
78    }
79
80    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
81        self.document_variable_name = name.into();
82        self
83    }
84
85    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
86        self.input_key = key.into();
87        self
88    }
89
90    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
91        self.output_key = key.into();
92        self
93    }
94
95    pub fn with_name(mut self, name: impl Into<String>) -> Self {
96        self.name = name.into();
97        self
98    }
99
100    pub fn with_verbose(mut self, verbose: bool) -> Self {
101        self.verbose = verbose;
102        self
103    }
104
105    pub fn build_initial_prompt(&self, context: &str, input: &str) -> String {
106        self.initial_prompt_template
107            .replace(&format!("{{{}}}", self.document_variable_name), context)
108            .replace("{input}", input)
109    }
110
111    pub fn build_refine_prompt(&self, context: &str, input: &str, existing_answer: &str) -> String {
112        self.refine_prompt_template
113            .replace(&format!("{{{}}}", self.document_variable_name), context)
114            .replace("{input}", input)
115            .replace("{existing_answer}", existing_answer)
116    }
117
118    /// Invoke with documents and input directly (iterative refinement).
119    pub async fn invoke_with_documents(
120        &self,
121        documents: Vec<Document>,
122        input: &str,
123    ) -> Result<String, ChainError>
124    where
125        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
126    {
127        if documents.is_empty() {
128            return Err(ChainError::ExecutionError(
129                "Document list is empty".to_string(),
130            ));
131        }
132
133        if self.verbose {
134            println!("\n=== RefineDocumentsChain ===");
135            println!("Document count: {}", documents.len());
136            println!("Input: {}", input);
137        }
138
139        // Step 1: Generate initial answer from the first document
140        let first_context = &documents[0].content;
141        let initial_prompt = self.build_initial_prompt(first_context, input);
142
143        if self.verbose {
144            println!("\n--- Initial processing (document 1) ---");
145        }
146
147        let messages = vec![Message::human(&initial_prompt)];
148        let response =
149            self.llm.invoke(messages, None).await.map_err(|e| {
150                ChainError::ExecutionError(format!("LLM initial call failed: {}", e))
151            })?;
152        let mut answer = response.content;
153
154        if self.verbose {
155            println!("Initial answer: {}", answer);
156        }
157
158        // Subsequent steps: iteratively refine with remaining documents
159        for (i, doc) in documents[1..].iter().enumerate() {
160            if self.verbose {
161                println!("\n--- Refinement step {} (document {}) ---", i + 1, i + 2);
162            }
163
164            let refine_prompt = self.build_refine_prompt(&doc.content, input, &answer);
165
166            let messages = vec![Message::human(&refine_prompt)];
167            let response = self.llm.invoke(messages, None).await.map_err(|e| {
168                ChainError::ExecutionError(format!("LLM refinement call failed: {}", e))
169            })?;
170            answer = response.content;
171
172            if self.verbose {
173                println!("Refined answer: {}", answer);
174            }
175        }
176
177        if self.verbose {
178            println!("=== RefineDocumentsChain complete ===\n");
179        }
180
181        Ok(answer)
182    }
183}
184
185#[async_trait]
186impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for RefineDocumentsChain<M>
187where
188    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
189{
190    fn input_keys(&self) -> Vec<&str> {
191        vec![&self.input_key, "documents"]
192    }
193
194    fn output_keys(&self) -> Vec<&str> {
195        vec![&self.output_key]
196    }
197
198    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
199        // P2-8: validate inputs on the invoke path too, matching stream.
200        self.validate_inputs(&inputs)?;
201
202        let input = inputs
203            .get(&self.input_key)
204            .and_then(|v| v.as_str())
205            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
206
207        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
208
209        let output = self.invoke_with_documents(documents, input).await?;
210
211        let mut result = HashMap::new();
212        result.insert(self.output_key.clone(), Value::String(output));
213        Ok(result)
214    }
215
216    /// Stream execution for RefineDocumentsChain.
217    ///
218    /// Runs the initial + all intermediate refine steps via invoke (since
219    /// their output feeds the next step), then streams the final refine step
220    /// token by token via `stream_chat`. With a single document there is no
221    /// final refine step — the initial answer is emitted directly (P2-4), so
222    /// the LLM is not re-called on the identical initial prompt.
223    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
224        self.validate_inputs(&inputs)?;
225
226        let input = inputs
227            .get(&self.input_key)
228            .and_then(|v| v.as_str())
229            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
230
231        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
232
233        if documents.is_empty() {
234            return Err(ChainError::ExecutionError(
235                "Document list is empty".to_string(),
236            ));
237        }
238
239        // Step 1: Generate initial answer from the first document
240        let first_context = &documents[0].content;
241        let initial_prompt = self.build_initial_prompt(first_context, input);
242        let messages = vec![Message::human(&initial_prompt)];
243        let response =
244            self.llm.invoke(messages, None).await.map_err(|e| {
245                ChainError::ExecutionError(format!("LLM initial call failed: {}", e))
246            })?;
247        let mut answer = response.content;
248
249        // Step 2: Run intermediate refine steps (all but the last) via invoke
250        //
251        // P2-4: with a single document `last_idx == 0` and `documents[1..0]`
252        // would panic on the slice index — iterate with skip/take so zero
253        // intermediate documents (1 or 2 documents) is a no-op and the initial
254        // answer flows straight to the final step.
255        let last_idx = documents.len() - 1;
256        for (i, doc) in documents
257            .iter()
258            .skip(1)
259            .take(last_idx.saturating_sub(1))
260            .enumerate()
261        {
262            let refine_prompt = self.build_refine_prompt(&doc.content, input, &answer);
263            let messages = vec![Message::human(&refine_prompt)];
264            let response = self.llm.invoke(messages, None).await.map_err(|e| {
265                ChainError::ExecutionError(format!("LLM refinement call failed: {}", e))
266            })?;
267            answer = response.content;
268
269            if self.verbose {
270                println!("Refine step {} completed", i + 1);
271            }
272        }
273
274        // Step 3: Stream the final refine step
275        //
276        // P2-4: with a single document the initial invoke above already produced
277        // the complete answer — calling `stream_chat` on the identical initial
278        // prompt would make a second, redundant LLM call for the same output
279        // (the invoke result was previously discarded and the prompt re-sent).
280        // Stream the computed answer directly instead.
281        if last_idx == 0 {
282            let stream = futures_util::stream::once(async move {
283                Ok(StreamToken {
284                    token: answer,
285                    is_final: true,
286                })
287            });
288            return Ok(Box::pin(stream));
289        }
290
291        let final_prompt = self.build_refine_prompt(&documents[last_idx].content, input, &answer);
292
293        let messages = vec![Message::human(&final_prompt)];
294        let llm_stream = self
295            .llm
296            .stream_chat(messages, None)
297            .await
298            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
299
300        let stream = llm_stream.map(|result| match result {
301            Ok(token) => Ok(StreamToken {
302                token,
303                is_final: false,
304            }),
305            Err(e) => Err(ChainError::StreamError(format!(
306                "Stream token error: {}",
307                e
308            ))),
309        });
310
311        let final_stream = stream.chain(futures_util::stream::once(async move {
312            Ok(StreamToken {
313                token: String::new(),
314                is_final: true,
315            })
316        }));
317
318        Ok(Box::pin(final_stream))
319    }
320
321    fn name(&self) -> &str {
322        &self.name
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use async_trait::async_trait;
330    use futures_util::Stream;
331    use lc_core::language_models::LLMResult;
332    use lc_core::runnables::RunnableConfig;
333    use lc_core::{BaseLanguageModel, Runnable};
334    use std::pin::Pin;
335    use std::sync::atomic::{AtomicUsize, Ordering};
336    use std::sync::Arc;
337
338    /// Mock chat model that counts `invoke`/`stream_chat` calls so the P2-4
339    /// single-document fix (no second LLM call) is provable.
340    #[derive(Debug)]
341    struct MockError(String);
342    impl std::fmt::Display for MockError {
343        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
344            write!(f, "{}", self.0)
345        }
346    }
347    impl std::error::Error for MockError {}
348
349    struct CountingLLM {
350        invokes: Arc<AtomicUsize>,
351        streams: Arc<AtomicUsize>,
352    }
353
354    #[async_trait]
355    impl Runnable<Vec<Message>, LLMResult> for CountingLLM {
356        type Error = MockError;
357        async fn invoke(
358            &self,
359            _input: Vec<Message>,
360            _config: Option<RunnableConfig>,
361        ) -> Result<LLMResult, Self::Error> {
362            self.invokes.fetch_add(1, Ordering::SeqCst);
363            Ok(LLMResult {
364                content: "initial answer".to_string(),
365                model: "mock".to_string(),
366                token_usage: None,
367                tool_calls: None,
368                thinking_content: None,
369            })
370        }
371    }
372
373    #[async_trait]
374    impl BaseLanguageModel<Vec<Message>, LLMResult> for CountingLLM {
375        fn model_name(&self) -> &str {
376            "mock"
377        }
378        fn get_num_tokens(&self, t: &str) -> usize {
379            t.len()
380        }
381        fn with_temperature(self, _: f32) -> Self {
382            self
383        }
384        fn with_max_tokens(self, _: usize) -> Self {
385            self
386        }
387    }
388
389    #[async_trait]
390    impl BaseChatModel for CountingLLM {
391        async fn chat(
392            &self,
393            _messages: Vec<Message>,
394            _config: Option<RunnableConfig>,
395        ) -> Result<LLMResult, Self::Error> {
396            self.invokes.fetch_add(1, Ordering::SeqCst);
397            Ok(LLMResult {
398                content: "initial answer".to_string(),
399                model: "mock".to_string(),
400                token_usage: None,
401                tool_calls: None,
402                thinking_content: None,
403            })
404        }
405        async fn stream_chat(
406            &self,
407            _messages: Vec<Message>,
408            _config: Option<RunnableConfig>,
409        ) -> Result<Pin<Box<dyn Stream<Item = Result<String, Self::Error>> + Send>>, Self::Error>
410        {
411            self.streams.fetch_add(1, Ordering::SeqCst);
412            let tokens = [Ok("refined answer".to_string())];
413            Ok(Box::pin(futures_util::stream::iter(tokens)))
414        }
415    }
416
417    fn inputs_for(documents: Vec<Document>) -> HashMap<String, Value> {
418        let mut inputs = HashMap::new();
419        inputs.insert("input".to_string(), Value::String("question".to_string()));
420        inputs.insert(
421            "documents".to_string(),
422            serde_json::to_value(documents).unwrap(),
423        );
424        inputs
425    }
426
427    /// P2-4: a single document reuses the invoke-computed initial answer —
428    /// `stream_chat` is never called on the identical prompt (one LLM call
429    /// total, not two).
430    #[tokio::test]
431    async fn test_refine_stream_single_document_skips_second_llm_call() {
432        let invokes = Arc::new(AtomicUsize::new(0));
433        let streams = Arc::new(AtomicUsize::new(0));
434        let chain = RefineDocumentsChain::new(CountingLLM {
435            invokes: invokes.clone(),
436            streams: streams.clone(),
437        });
438        let inputs = inputs_for(vec![Document::new("doc one")]);
439
440        let mut stream = chain.stream(inputs).await.unwrap();
441        let mut tokens = Vec::new();
442        while let Some(item) = stream.next().await {
443            tokens.push(item.unwrap());
444        }
445        let text: String = tokens.iter().map(|t| t.token.as_str()).collect();
446        assert_eq!(text, "initial answer");
447        assert!(tokens.last().unwrap().is_final);
448        assert_eq!(invokes.load(Ordering::SeqCst), 1, "one initial invoke");
449        assert_eq!(
450            streams.load(Ordering::SeqCst),
451            0,
452            "single-document stream must not re-call the LLM"
453        );
454    }
455
456    /// Multi-document: initial + intermediate refines run via invoke, the final
457    /// refine is genuinely streamed (one `stream_chat` call).
458    #[tokio::test]
459    async fn test_refine_stream_multi_document_streams_final_refine() {
460        let invokes = Arc::new(AtomicUsize::new(0));
461        let streams = Arc::new(AtomicUsize::new(0));
462        let chain = RefineDocumentsChain::new(CountingLLM {
463            invokes: invokes.clone(),
464            streams: streams.clone(),
465        });
466        let inputs = inputs_for(vec![
467            Document::new("doc one"),
468            Document::new("doc two"),
469            Document::new("doc three"),
470        ]);
471
472        let mut stream = chain.stream(inputs).await.unwrap();
473        let mut tokens = Vec::new();
474        while let Some(item) = stream.next().await {
475            tokens.push(item.unwrap());
476        }
477        let text: String = tokens.iter().map(|t| t.token.as_str()).collect();
478        assert!(text.contains("refined answer"));
479        assert!(tokens.last().unwrap().is_final);
480        // 3 documents → 1 initial + 1 intermediate invoke, then 1 streamed final.
481        assert_eq!(invokes.load(Ordering::SeqCst), 2);
482        assert_eq!(streams.load(Ordering::SeqCst), 1);
483    }
484
485    #[tokio::test]
486    async fn test_refine_stream_empty_documents() {
487        let chain = RefineDocumentsChain::new(CountingLLM {
488            invokes: Arc::new(AtomicUsize::new(0)),
489            streams: Arc::new(AtomicUsize::new(0)),
490        });
491        let mut inputs = HashMap::new();
492        inputs.insert("input".to_string(), Value::String("q".to_string()));
493        inputs.insert("documents".to_string(), serde_json::json!([]));
494        let err = match chain.stream(inputs).await {
495            Ok(_) => panic!("expected an execution error"),
496            Err(e) => e,
497        };
498        assert!(matches!(err, ChainError::ExecutionError(_)));
499    }
500}
lc_chains/document_chains/refine.rs

lc_chains/document_chains/
refine.rs
