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::BaseChatModel;
7use lc_providers::{wrap_chat_model, ProviderError};
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};
14use crate::BoxedChatModel;
15
16/// Default initial processing prompt template.
17pub(crate) const DEFAULT_REFINE_INITIAL_PROMPT: &str =
18    "Answer the question based on the following reference information.
19
20Reference information:
21{context}
22
23Question: {input}
24
25Answer:";
26
27/// Default iterative refinement prompt template.
28pub(crate) const DEFAULT_REFINE_PROMPT: &str = "You have provided an answer based on partial information. Here is additional reference information.
29
30Existing answer:
31{existing_answer}
32
33New reference information:
34{context}
35
36Please 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.
37
38Question: {input}
39
40Refined answer:";
41
42/// RefineDocumentsChain
43///
44/// Iteratively refines the answer document by document.
45/// Generates an initial answer from the first document, then refines with each subsequent document.
46pub struct RefineDocumentsChain {
47    llm: BoxedChatModel,
48    initial_prompt_template: String,
49    refine_prompt_template: String,
50    document_variable_name: String,
51    input_key: String,
52    output_key: String,
53    name: String,
54    verbose: bool,
55}
56
57impl RefineDocumentsChain {
58    /// Create a new [`RefineDocumentsChain`] with the given LLM.
59    pub fn new<L>(llm: L) -> Self
60    where
61        L: BaseChatModel + Send + Sync + 'static,
62        L::Error: Into<ProviderError>,
63    {
64        Self {
65            llm: wrap_chat_model(llm),
66            initial_prompt_template: DEFAULT_REFINE_INITIAL_PROMPT.to_string(),
67            refine_prompt_template: DEFAULT_REFINE_PROMPT.to_string(),
68            document_variable_name: "context".to_string(),
69            input_key: "input".to_string(),
70            output_key: "output".to_string(),
71            name: "refine_documents".to_string(),
72            verbose: false,
73        }
74    }
75
76    /// Set the initial prompt template used for the first document.
77    pub fn with_initial_prompt(mut self, template: impl Into<String>) -> Self {
78        self.initial_prompt_template = template.into();
79        self
80    }
81
82    /// Set the refine prompt template used for subsequent documents.
83    pub fn with_refine_prompt(mut self, template: impl Into<String>) -> Self {
84        self.refine_prompt_template = template.into();
85        self
86    }
87
88    /// Set the document variable name used in the prompts.
89    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
90        self.document_variable_name = name.into();
91        self
92    }
93
94    /// Set the input key.
95    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
96        self.input_key = key.into();
97        self
98    }
99
100    /// Set the output key.
101    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
102        self.output_key = key.into();
103        self
104    }
105
106    /// Set the chain name.
107    pub fn with_name(mut self, name: impl Into<String>) -> Self {
108        self.name = name.into();
109        self
110    }
111
112    /// Set verbose mode.
113    pub fn with_verbose(mut self, verbose: bool) -> Self {
114        self.verbose = verbose;
115        self
116    }
117
118    /// Build the initial prompt from the first document's content and input.
119    pub fn build_initial_prompt(&self, context: &str, input: &str) -> String {
120        self.initial_prompt_template
121            .replace(&format!("{{{}}}", self.document_variable_name), context)
122            .replace("{input}", input)
123    }
124
125    /// Build the refine prompt from the new context, input, and existing answer.
126    pub fn build_refine_prompt(&self, context: &str, input: &str, existing_answer: &str) -> String {
127        self.refine_prompt_template
128            .replace(&format!("{{{}}}", self.document_variable_name), context)
129            .replace("{input}", input)
130            .replace("{existing_answer}", existing_answer)
131    }
132
133    /// Invoke with documents and input directly (iterative refinement).
134    pub async fn invoke_with_documents(
135        &self,
136        documents: Vec<Document>,
137        input: &str,
138    ) -> Result<String, ChainError> {
139        if documents.is_empty() {
140            return Err(ChainError::ExecutionError(
141                "Document list is empty".to_string(),
142            ));
143        }
144
145        if self.verbose {
146            println!("\n=== RefineDocumentsChain ===");
147            println!("Document count: {}", documents.len());
148            println!("Input: {}", input);
149        }
150
151        // Step 1: Generate initial answer from the first document
152        let first_context = &documents[0].content;
153        let initial_prompt = self.build_initial_prompt(first_context, input);
154
155        if self.verbose {
156            println!("\n--- Initial processing (document 1) ---");
157        }
158
159        let messages = vec![Message::human(&initial_prompt)];
160        let response =
161            self.llm.invoke(messages, None).await.map_err(|e| {
162                ChainError::ExecutionError(format!("LLM initial call failed: {}", e))
163            })?;
164        let mut answer = response.content;
165
166        if self.verbose {
167            println!("Initial answer: {}", answer);
168        }
169
170        // Subsequent steps: iteratively refine with remaining documents
171        for (i, doc) in documents[1..].iter().enumerate() {
172            if self.verbose {
173                println!("\n--- Refinement step {} (document {}) ---", i + 1, i + 2);
174            }
175
176            let refine_prompt = self.build_refine_prompt(&doc.content, input, &answer);
177
178            let messages = vec![Message::human(&refine_prompt)];
179            let response = self.llm.invoke(messages, None).await.map_err(|e| {
180                ChainError::ExecutionError(format!("LLM refinement call failed: {}", e))
181            })?;
182            answer = response.content;
183
184            if self.verbose {
185                println!("Refined answer: {}", answer);
186            }
187        }
188
189        if self.verbose {
190            println!("=== RefineDocumentsChain complete ===\n");
191        }
192
193        Ok(answer)
194    }
195}
196
197#[async_trait]
198impl BaseChain for RefineDocumentsChain {
199    fn input_keys(&self) -> Vec<&str> {
200        vec![&self.input_key, "documents"]
201    }
202
203    fn output_keys(&self) -> Vec<&str> {
204        vec![&self.output_key]
205    }
206
207    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
208        // P2-8: validate inputs on the invoke path too, matching stream.
209        self.validate_inputs(&inputs)?;
210
211        let input = inputs
212            .get(&self.input_key)
213            .and_then(|v| v.as_str())
214            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
215
216        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
217
218        let output = self.invoke_with_documents(documents, input).await?;
219
220        let mut result = HashMap::new();
221        result.insert(self.output_key.clone(), Value::String(output));
222        Ok(result)
223    }
224
225    /// Stream execution for RefineDocumentsChain.
226    ///
227    /// Runs the initial + all intermediate refine steps via invoke (since
228    /// their output feeds the next step), then streams the final refine step
229    /// token by token via `stream_chat`. With a single document there is no
230    /// final refine step — the initial answer is emitted directly (P2-4), so
231    /// the LLM is not re-called on the identical initial prompt.
232    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
233        self.validate_inputs(&inputs)?;
234
235        let input = inputs
236            .get(&self.input_key)
237            .and_then(|v| v.as_str())
238            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
239
240        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
241
242        if documents.is_empty() {
243            return Err(ChainError::ExecutionError(
244                "Document list is empty".to_string(),
245            ));
246        }
247
248        // Step 1: Generate initial answer from the first document
249        let first_context = &documents[0].content;
250        let initial_prompt = self.build_initial_prompt(first_context, input);
251        let messages = vec![Message::human(&initial_prompt)];
252        let response =
253            self.llm.invoke(messages, None).await.map_err(|e| {
254                ChainError::ExecutionError(format!("LLM initial call failed: {}", e))
255            })?;
256        let mut answer = response.content;
257
258        // Step 2: Run intermediate refine steps (all but the last) via invoke
259        //
260        // P2-4: with a single document `last_idx == 0` and `documents[1..0]`
261        // would panic on the slice index — iterate with skip/take so zero
262        // intermediate documents (1 or 2 documents) is a no-op and the initial
263        // answer flows straight to the final step.
264        let last_idx = documents.len() - 1;
265        for (i, doc) in documents
266            .iter()
267            .skip(1)
268            .take(last_idx.saturating_sub(1))
269            .enumerate()
270        {
271            let refine_prompt = self.build_refine_prompt(&doc.content, input, &answer);
272            let messages = vec![Message::human(&refine_prompt)];
273            let response = self.llm.invoke(messages, None).await.map_err(|e| {
274                ChainError::ExecutionError(format!("LLM refinement call failed: {}", e))
275            })?;
276            answer = response.content;
277
278            if self.verbose {
279                println!("Refine step {} completed", i + 1);
280            }
281        }
282
283        // Step 3: Stream the final refine step
284        //
285        // P2-4: with a single document the initial invoke above already produced
286        // the complete answer — calling `stream_chat` on the identical initial
287        // prompt would make a second, redundant LLM call for the same output
288        // (the invoke result was previously discarded and the prompt re-sent).
289        // Stream the computed answer directly instead.
290        if last_idx == 0 {
291            let stream = futures_util::stream::once(async move {
292                Ok(StreamToken {
293                    token: answer,
294                    is_final: true,
295                })
296            });
297            return Ok(Box::pin(stream));
298        }
299
300        let final_prompt = self.build_refine_prompt(&documents[last_idx].content, input, &answer);
301
302        let messages = vec![Message::human(&final_prompt)];
303        let llm_stream = self
304            .llm
305            .stream_chat(messages, None)
306            .await
307            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
308
309        let stream = llm_stream.map(|result| match result {
310            Ok(token) => Ok(StreamToken {
311                token,
312                is_final: false,
313            }),
314            Err(e) => Err(ChainError::StreamError(format!(
315                "Stream token error: {}",
316                e
317            ))),
318        });
319
320        let final_stream = stream.chain(futures_util::stream::once(async move {
321            Ok(StreamToken {
322                token: String::new(),
323                is_final: true,
324            })
325        }));
326
327        Ok(Box::pin(final_stream))
328    }
329
330    fn name(&self) -> &str {
331        &self.name
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338    use async_trait::async_trait;
339    use futures_util::Stream;
340    use lc_core::language_models::LLMResult;
341    use lc_core::runnables::RunnableConfig;
342    use lc_core::{BaseLanguageModel, Runnable};
343    use std::pin::Pin;
344    use std::sync::atomic::{AtomicUsize, Ordering};
345    use std::sync::Arc;
346
347    /// Mock chat model that counts `invoke`/`stream_chat` calls so the P2-4
348    /// single-document fix (no second LLM call) is provable.
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 = ProviderError;
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
