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 lc_core::language_models::LLMResult;
6use lc_core::{BaseChatModel, Runnable};
7use lc_schema::Message;
8use lc_shared::document::Document;
9use serde_json::Value;
10use std::collections::HashMap;
11
12use crate::base::{BaseChain, ChainError, ChainResult};
13
14/// Default initial processing prompt template.
15pub(crate) const DEFAULT_REFINE_INITIAL_PROMPT: &str =
16    "Answer the question based on the following reference information.
17
18Reference information:
19{context}
20
21Question: {input}
22
23Answer:";
24
25/// Default iterative refinement prompt template.
26pub(crate) const DEFAULT_REFINE_PROMPT: &str = "You have provided an answer based on partial information. Here is additional reference information.
27
28Existing answer:
29{existing_answer}
30
31New reference information:
32{context}
33
34Please 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.
35
36Question: {input}
37
38Refined answer:";
39
40/// RefineDocumentsChain
41///
42/// Iteratively refines the answer document by document.
43/// Generates an initial answer from the first document, then refines with each subsequent document.
44pub struct RefineDocumentsChain<M: BaseChatModel> {
45    llm: M,
46    initial_prompt_template: String,
47    refine_prompt_template: String,
48    document_variable_name: String,
49    input_key: String,
50    output_key: String,
51    name: String,
52    verbose: bool,
53}
54
55impl<M: BaseChatModel> RefineDocumentsChain<M> {
56    pub fn new(llm: M) -> Self {
57        Self {
58            llm,
59            initial_prompt_template: DEFAULT_REFINE_INITIAL_PROMPT.to_string(),
60            refine_prompt_template: DEFAULT_REFINE_PROMPT.to_string(),
61            document_variable_name: "context".to_string(),
62            input_key: "input".to_string(),
63            output_key: "output".to_string(),
64            name: "refine_documents".to_string(),
65            verbose: false,
66        }
67    }
68
69    pub fn with_initial_prompt(mut self, template: impl Into<String>) -> Self {
70        self.initial_prompt_template = template.into();
71        self
72    }
73
74    pub fn with_refine_prompt(mut self, template: impl Into<String>) -> Self {
75        self.refine_prompt_template = template.into();
76        self
77    }
78
79    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
80        self.document_variable_name = name.into();
81        self
82    }
83
84    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
85        self.input_key = key.into();
86        self
87    }
88
89    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
90        self.output_key = key.into();
91        self
92    }
93
94    pub fn with_name(mut self, name: impl Into<String>) -> Self {
95        self.name = name.into();
96        self
97    }
98
99    pub fn with_verbose(mut self, verbose: bool) -> Self {
100        self.verbose = verbose;
101        self
102    }
103
104    pub fn build_initial_prompt(&self, context: &str, input: &str) -> String {
105        self.initial_prompt_template
106            .replace(&format!("{{{}}}", self.document_variable_name), context)
107            .replace("{input}", input)
108    }
109
110    pub fn build_refine_prompt(&self, context: &str, input: &str, existing_answer: &str) -> String {
111        self.refine_prompt_template
112            .replace(&format!("{{{}}}", self.document_variable_name), context)
113            .replace("{input}", input)
114            .replace("{existing_answer}", existing_answer)
115    }
116
117    /// Invoke with documents and input directly (iterative refinement).
118    pub async fn invoke_with_documents(
119        &self,
120        documents: Vec<Document>,
121        input: &str,
122    ) -> Result<String, ChainError>
123    where
124        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
125    {
126        if documents.is_empty() {
127            return Err(ChainError::ExecutionError(
128                "Document list is empty".to_string(),
129            ));
130        }
131
132        if self.verbose {
133            println!("\n=== RefineDocumentsChain ===");
134            println!("Document count: {}", documents.len());
135            println!("Input: {}", input);
136        }
137
138        // Step 1: Generate initial answer from the first document
139        let first_context = &documents[0].content;
140        let initial_prompt = self.build_initial_prompt(first_context, input);
141
142        if self.verbose {
143            println!("\n--- Initial processing (document 1) ---");
144        }
145
146        let messages = vec![Message::human(&initial_prompt)];
147        let response =
148            self.llm.invoke(messages, None).await.map_err(|e| {
149                ChainError::ExecutionError(format!("LLM initial call failed: {}", e))
150            })?;
151        let mut answer = response.content;
152
153        if self.verbose {
154            println!("Initial answer: {}", answer);
155        }
156
157        // Subsequent steps: iteratively refine with remaining documents
158        for (i, doc) in documents[1..].iter().enumerate() {
159            if self.verbose {
160                println!("\n--- Refinement step {} (document {}) ---", i + 1, i + 2);
161            }
162
163            let refine_prompt = self.build_refine_prompt(&doc.content, input, &answer);
164
165            let messages = vec![Message::human(&refine_prompt)];
166            let response = self.llm.invoke(messages, None).await.map_err(|e| {
167                ChainError::ExecutionError(format!("LLM refinement call failed: {}", e))
168            })?;
169            answer = response.content;
170
171            if self.verbose {
172                println!("Refined answer: {}", answer);
173            }
174        }
175
176        if self.verbose {
177            println!("=== RefineDocumentsChain complete ===\n");
178        }
179
180        Ok(answer)
181    }
182}
183
184#[async_trait]
185impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for RefineDocumentsChain<M>
186where
187    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
188{
189    fn input_keys(&self) -> Vec<&str> {
190        vec![&self.input_key, "documents"]
191    }
192
193    fn output_keys(&self) -> Vec<&str> {
194        vec![&self.output_key]
195    }
196
197    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
198        let input = inputs
199            .get(&self.input_key)
200            .and_then(|v| v.as_str())
201            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
202
203        let documents: Vec<Document> = inputs
204            .get("documents")
205            .and_then(|v| v.as_array())
206            .map(|arr| {
207                arr.iter()
208                    .filter_map(|v| serde_json::from_value(v.clone()).ok())
209                    .collect()
210            })
211            .ok_or_else(|| ChainError::MissingInput("documents".to_string()))?;
212
213        let output = self.invoke_with_documents(documents, input).await?;
214
215        let mut result = HashMap::new();
216        result.insert(self.output_key.clone(), Value::String(output));
217        Ok(result)
218    }
219
220    fn name(&self) -> &str {
221        &self.name
222    }
223}
lc_chains/document_chains/refine.rs

lc_chains/document_chains/
refine.rs
