lc_rag/
multi_query.rs

1// src/retrieval/multi_query.rs
2//! MultiQueryRetriever 实现
3//!
4//! 使用 LLM 生成多个查询变体，提高检索召回率。
5
6use lc_core::Runnable;
7use lc_prompts::PromptTemplate;
8use lc_providers::OpenAIChat;
9use lc_schema::Message;
10use lc_vector_stores::{Document, SearchResult};
11
12use crate::retriever::RetrieverTrait;
13use std::collections::HashMap;
14use std::sync::Arc;
15
16/// Generate a stable document ID from content hash (M58).
17fn doc_content_hash(content: &str) -> String {
18    use std::hash::{Hash, Hasher};
19    let mut hasher = std::collections::hash_map::DefaultHasher::new();
20    content.hash(&mut hasher);
21    format!("{:016x}", hasher.finish())
22}
23
24/// MultiQueryRetriever 错误类型
25#[derive(Debug)]
26pub enum MultiQueryError {
27    /// LLM 错误
28    LLMError(String),
29
30    /// 检索错误
31    RetrieverError(String),
32
33    /// 解析错误
34    ParseError(String),
35}
36
37impl std::fmt::Display for MultiQueryError {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        match self {
40            MultiQueryError::LLMError(msg) => write!(f, "LLM 错误: {}", msg),
41            MultiQueryError::RetrieverError(msg) => write!(f, "检索错误: {}", msg),
42            MultiQueryError::ParseError(msg) => write!(f, "解析错误: {}", msg),
43        }
44    }
45}
46
47impl std::error::Error for MultiQueryError {}
48
49/// MultiQueryRetriever 配置
50pub struct MultiQueryConfig {
51    /// 生成的查询数量
52    pub num_queries: usize,
53
54    /// 每个查询返回的文档数
55    pub k_per_query: usize,
56
57    /// 最终返回的文档数
58    pub final_k: usize,
59
60    /// 查询生成 prompt
61    pub prompt_template: String,
62}
63
64impl Default for MultiQueryConfig {
65    fn default() -> Self {
66        Self {
67            num_queries: 3,
68            k_per_query: 5,
69            final_k: 10,
70            prompt_template: DEFAULT_MULTI_QUERY_PROMPT.to_string(),
71        }
72    }
73}
74
75impl MultiQueryConfig {
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    pub fn with_num_queries(mut self, n: usize) -> Self {
81        self.num_queries = n;
82        self
83    }
84
85    pub fn with_k_per_query(mut self, k: usize) -> Self {
86        self.k_per_query = k;
87        self
88    }
89
90    pub fn with_final_k(mut self, k: usize) -> Self {
91        self.final_k = k;
92        self
93    }
94
95    pub fn with_prompt(mut self, prompt: String) -> Self {
96        self.prompt_template = prompt;
97        self
98    }
99}
100
101const DEFAULT_MULTI_QUERY_PROMPT: &str = r#"You are an AI language model assistant. Your task is to generate 3 different versions of the given user question to retrieve relevant documents from a vector database.
102
103By generating multiple perspectives on the user question, your goal is to help overcome some of the limitations of distance-based similarity search.
104
105Provide these alternative questions separated by newlines.
106
107Original question: {question}
108
109Alternative questions:"#;
110
111/// MultiQueryRetriever
112///
113/// 使用 LLM 生成多个查询变体，然后用基础检索器分别检索，
114/// 最后合并去重结果返回。
115pub struct MultiQueryRetriever {
116    /// LLM 用于生成查询变体
117    llm: OpenAIChat,
118
119    /// 基础检索器
120    base_retriever: Arc<dyn RetrieverTrait>,
121
122    /// 配置
123    config: MultiQueryConfig,
124}
125
126impl MultiQueryRetriever {
127    pub fn new(llm: OpenAIChat, base_retriever: Arc<dyn RetrieverTrait>) -> Self {
128        Self {
129            llm,
130            base_retriever,
131            config: MultiQueryConfig::default(),
132        }
133    }
134
135    pub fn with_config(mut self, config: MultiQueryConfig) -> Self {
136        self.config = config;
137        self
138    }
139
140    pub fn with_num_queries(mut self, n: usize) -> Self {
141        self.config.num_queries = n;
142        self
143    }
144
145    pub fn with_k_per_query(mut self, k: usize) -> Self {
146        self.config.k_per_query = k;
147        self
148    }
149
150    pub fn with_final_k(mut self, k: usize) -> Self {
151        self.config.final_k = k;
152        self
153    }
154
155    async fn generate_queries(&self, original_query: &str) -> Result<Vec<String>, MultiQueryError> {
156        let template = PromptTemplate::new(&self.config.prompt_template);
157        let mut vars = HashMap::new();
158        vars.insert("question", original_query);
159        let prompt = template
160            .format(&vars)
161            .unwrap_or_else(|_| self.config.prompt_template.clone());
162
163        let messages = vec![Message::human(prompt)];
164
165        let response = self
166            .llm
167            .invoke(messages, None)
168            .await
169            .map_err(|e| MultiQueryError::LLMError(e.to_string()))?;
170
171        let content = response.content;
172
173        let queries: Vec<String> = content
174            .lines()
175            .filter(|line| !line.trim().is_empty())
176            .map(|line| line.trim().to_string())
177            .collect();
178
179        if queries.is_empty() {
180            return Err(MultiQueryError::ParseError(
181                "LLM 未生成有效的查询变体".to_string(),
182            ));
183        }
184
185        Ok(queries)
186    }
187
188    pub async fn retrieve_multi(&self, query: &str) -> Result<Vec<Document>, MultiQueryError> {
189        let queries = self.generate_queries(query).await?;
190
191        let all_queries: Vec<String> = std::iter::once(query.to_string()).chain(queries).collect();
192
193        let mut doc_scores: HashMap<String, (Document, f32)> = HashMap::new();
194
195        for q in &all_queries {
196            let results = self
197                .base_retriever
198                .retrieve_with_scores(q, self.config.k_per_query)
199                .await
200                .map_err(|e| MultiQueryError::RetrieverError(e.to_string()))?;
201
202            for result in results {
203                let doc_id = result
204                    .document
205                    .id
206                    .clone()
207                    .unwrap_or_else(|| doc_content_hash(&result.document.content));
208
209                doc_scores
210                    .entry(doc_id)
211                    .and_modify(|(_, score)| {
212                        // M3: use addition instead of no-op .max()
213                        *score += result.score;
214                    })
215                    .or_insert((result.document.clone(), result.score));
216            }
217        }
218
219        let mut scored_docs: Vec<(Document, f32)> = doc_scores
220            .values()
221            .map(|(doc, score)| (doc.clone(), *score))
222            .collect();
223
224        scored_docs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
225
226        let final_docs: Vec<Document> = scored_docs
227            .into_iter()
228            .take(self.config.final_k)
229            .map(|(doc, _)| doc)
230            .collect();
231
232        Ok(final_docs)
233    }
234
235    pub async fn retrieve_multi_with_scores(
236        &self,
237        query: &str,
238    ) -> Result<Vec<SearchResult>, MultiQueryError> {
239        let queries = self.generate_queries(query).await?;
240
241        let all_queries: Vec<String> = std::iter::once(query.to_string()).chain(queries).collect();
242
243        let mut doc_scores: HashMap<String, (Document, f32, usize)> = HashMap::new();
244
245        for q in &all_queries {
246            let results = self
247                .base_retriever
248                .retrieve_with_scores(q, self.config.k_per_query)
249                .await
250                .map_err(|e| MultiQueryError::RetrieverError(e.to_string()))?;
251
252            for result in results {
253                let doc_id = result
254                    .document
255                    .id
256                    .clone()
257                    .unwrap_or_else(|| doc_content_hash(&result.document.content));
258
259                doc_scores
260                    .entry(doc_id)
261                    .and_modify(|(_, score, count)| {
262                        // M3: use addition instead of no-op .max()
263                        *score += result.score;
264                        *count += 1;
265                    })
266                    .or_insert((result.document.clone(), result.score, 1));
267            }
268        }
269
270        let mut scored_docs: Vec<SearchResult> = doc_scores
271            .values()
272            .map(|(doc, score, count)| {
273                let combined_score = score * (1.0 + 0.1 * *count as f32);
274                SearchResult {
275                    document: doc.clone(),
276                    score: combined_score,
277                }
278            })
279            .collect();
280
281        scored_docs.sort_by(|a, b| {
282            b.score
283                .partial_cmp(&a.score)
284                .unwrap_or(std::cmp::Ordering::Equal)
285        });
286
287        let final_results: Vec<SearchResult> =
288            scored_docs.into_iter().take(self.config.final_k).collect();
289
290        Ok(final_results)
291    }
292
293    pub async fn get_generated_queries(&self, query: &str) -> Result<Vec<String>, MultiQueryError> {
294        self.generate_queries(query).await
295    }
296}
297
298/// 静态查询生成器（不依赖 LLM）
299#[allow(clippy::type_complexity)]
300pub struct StaticQueryGenerator {
301    expansions: Vec<Box<dyn Fn(&str) -> Vec<String> + Send + Sync>>,
302}
303
304impl StaticQueryGenerator {
305    pub fn new() -> Self {
306        Self {
307            expansions: Vec::new(),
308        }
309    }
310
311    pub fn with_synonym_expansion(mut self, synonyms: HashMap<String, Vec<String>>) -> Self {
312        self.expansions.push(Box::new(move |query: &str| {
313            let mut expanded = Vec::new();
314            for (word, syns) in &synonyms {
315                if query.contains(word) {
316                    for syn in syns {
317                        expanded.push(query.replace(word, syn));
318                    }
319                }
320            }
321            expanded
322        }));
323        self
324    }
325
326    pub fn with_prefix_expansion(mut self, prefixes: Vec<String>) -> Self {
327        self.expansions.push(Box::new(move |query: &str| {
328            prefixes
329                .iter()
330                .map(|p| format!("{} {}", p, query))
331                .collect()
332        }));
333        self
334    }
335
336    pub fn generate(&self, query: &str) -> Vec<String> {
337        self.expansions
338            .iter()
339            .flat_map(|exp| exp(query))
340            .filter(|q| q != query)
341            .collect()
342    }
343}
344
345impl Default for StaticQueryGenerator {
346    fn default() -> Self {
347        Self::new()
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354
355    #[test]
356    fn test_static_query_generator_synonym() {
357        let synonyms: HashMap<String, Vec<String>> = HashMap::from([(
358            "数据库".to_string(),
359            vec!["DB".to_string(), "存储".to_string()],
360        )]);
361
362        let generator = StaticQueryGenerator::new().with_synonym_expansion(synonyms);
363
364        let queries = generator.generate("数据库连接失败");
365
366        assert!(queries.contains(&"DB连接失败".to_string()));
367        assert!(queries.contains(&"存储连接失败".to_string()));
368    }
369
370    #[test]
371    fn test_static_query_generator_prefix() {
372        let generator = StaticQueryGenerator::new()
373            .with_prefix_expansion(vec!["如何".to_string(), "怎么".to_string()]);
374
375        let queries = generator.generate("处理错误");
376
377        assert!(queries.contains(&"如何 处理错误".to_string()));
378        assert!(queries.contains(&"怎么 处理错误".to_string()));
379    }
380
381    #[test]
382    fn test_multi_query_config() {
383        let config = MultiQueryConfig::new()
384            .with_num_queries(5)
385            .with_k_per_query(10)
386            .with_final_k(20);
387
388        assert_eq!(config.num_queries, 5);
389        assert_eq!(config.k_per_query, 10);
390        assert_eq!(config.final_k, 20);
391    }
392
393    #[test]
394    fn test_multi_query_config_default() {
395        let config = MultiQueryConfig::default();
396
397        assert_eq!(config.num_queries, 3);
398        assert_eq!(config.k_per_query, 5);
399        assert_eq!(config.final_k, 10);
400    }
401}
lc_rag/multi_query.rs

lc_rag/
multi_query.rs
