lc_rag/
multi_query.rs

1// src/retrieval/multi_query.rs
2//! MultiQueryRetriever 实现
3//!
4//! 使用 LLM 生成多个查询变体，提高检索召回率。
5
6use lc_core::language_models::BaseChatModel;
7use lc_core::tools::ToolDefinition;
8use lc_prompts::PromptTemplate;
9use lc_providers::ProviderError;
10use lc_schema::Message;
11use lc_vector_stores::{Document, SearchResult};
12use serde_json::json;
13
14use crate::retriever::RetrieverTrait;
15use crate::structured::chat_structured;
16use std::collections::HashMap;
17use std::sync::Arc;
18
19/// Generate a stable document ID from content hash (M58).
20///
21/// P2-3: 用 FNV-1a 64 替代 `DefaultHasher`(std 内部算法不保证跨进程稳定;
22/// FNV-1a 是完全指定的确定性哈希)。
23fn doc_content_hash(content: &str) -> String {
24    use std::hash::{Hash, Hasher};
25    let mut hasher = fnv::FnvHasher::default();
26    content.hash(&mut hasher);
27    format!("{:016x}", hasher.finish())
28}
29
30/// MultiQueryRetriever 错误类型
31#[derive(Debug)]
32#[non_exhaustive]
33pub enum MultiQueryError {
34    /// LLM 错误
35    LLMError(String),
36
37    /// 检索错误
38    RetrieverError(String),
39
40    /// 解析错误
41    ParseError(String),
42}
43
44impl std::fmt::Display for MultiQueryError {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        match self {
47            MultiQueryError::LLMError(msg) => write!(f, "LLM error: {}", msg),
48            MultiQueryError::RetrieverError(msg) => write!(f, "Retriever error: {}", msg),
49            MultiQueryError::ParseError(msg) => write!(f, "Parse error: {}", msg),
50        }
51    }
52}
53
54impl std::error::Error for MultiQueryError {}
55
56/// MultiQueryRetriever 配置
57pub struct MultiQueryConfig {
58    /// 生成的查询数量
59    pub num_queries: usize,
60
61    /// 每个查询返回的文档数
62    pub k_per_query: usize,
63
64    /// 最终返回的文档数
65    pub final_k: usize,
66
67    /// 查询生成 prompt
68    pub prompt_template: String,
69}
70
71impl Default for MultiQueryConfig {
72    fn default() -> Self {
73        Self {
74            num_queries: 3,
75            k_per_query: 5,
76            final_k: 10,
77            prompt_template: DEFAULT_MULTI_QUERY_PROMPT.to_string(),
78        }
79    }
80}
81
82impl MultiQueryConfig {
83    /// 创建使用默认配置的 `MultiQueryConfig`
84    pub fn new() -> Self {
85        Self::default()
86    }
87
88    /// 设置生成的查询数量
89    pub fn with_num_queries(mut self, n: usize) -> Self {
90        self.num_queries = n;
91        self
92    }
93
94    /// 设置每个查询返回的文档数
95    pub fn with_k_per_query(mut self, k: usize) -> Self {
96        self.k_per_query = k;
97        self
98    }
99
100    /// 设置最终返回的文档数
101    pub fn with_final_k(mut self, k: usize) -> Self {
102        self.final_k = k;
103        self
104    }
105
106    /// 设置查询生成 prompt
107    pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
108        self.prompt_template = prompt.into();
109        self
110    }
111}
112
113const 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.
114
115By generating multiple perspectives on the user question, your goal is to help overcome some of the limitations of distance-based similarity search.
116
117Provide these alternative questions separated by newlines.
118
119Original question: {question}
120
121Alternative questions:"#;
122
123/// MultiQueryRetriever
124///
125/// 使用 LLM 生成多个查询变体，然后用基础检索器分别检索，
126/// 最后合并去重结果返回。
127pub struct MultiQueryRetriever {
128    /// LLM 用于生成查询变体
129    ///
130    /// P0-3: 不再硬编码 `OpenAIChat`,接受任意实现 `BaseChatModel` 的 LLM。
131    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
132
133    /// 基础检索器
134    base_retriever: Arc<dyn RetrieverTrait>,
135
136    /// 配置
137    config: MultiQueryConfig,
138}
139
140impl MultiQueryRetriever {
141    /// 创建 MultiQueryRetriever(接受任意实现 `BaseChatModel` 的 LLM)
142    pub fn new<L>(llm: L, base_retriever: Arc<dyn RetrieverTrait>) -> Self
143    where
144        L: BaseChatModel + Send + Sync + 'static,
145        L::Error: Into<ProviderError>,
146    {
147        Self {
148            llm: lc_providers::wrap_chat_model(llm),
149            base_retriever,
150            config: MultiQueryConfig::default(),
151        }
152    }
153
154    /// P0-3: 从已包装的 `Arc<dyn BaseChatModel<Error = ProviderError>>` 构建
155    pub fn new_arc(
156        llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
157        base_retriever: Arc<dyn RetrieverTrait>,
158    ) -> Self {
159        Self {
160            llm,
161            base_retriever,
162            config: MultiQueryConfig::default(),
163        }
164    }
165
166    /// 设置 MultiQuery 配置
167    pub fn with_config(mut self, config: MultiQueryConfig) -> Self {
168        self.config = config;
169        self
170    }
171
172    /// 设置生成的查询数量
173    pub fn with_num_queries(mut self, n: usize) -> Self {
174        self.config.num_queries = n;
175        self
176    }
177
178    /// 设置每个查询返回的文档数
179    pub fn with_k_per_query(mut self, k: usize) -> Self {
180        self.config.k_per_query = k;
181        self
182    }
183
184    /// 设置最终返回的文档数
185    pub fn with_final_k(mut self, k: usize) -> Self {
186        self.config.final_k = k;
187        self
188    }
189
190    async fn generate_queries(&self, original_query: &str) -> Result<Vec<String>, MultiQueryError> {
191        let template = PromptTemplate::new(&self.config.prompt_template);
192        let mut vars = HashMap::new();
193        vars.insert("question", original_query);
194        let prompt = template
195            .format(&vars)
196            .unwrap_or_else(|_| self.config.prompt_template.clone());
197
198        // P2-1: 优先 tool_calls 结构化查询列表;文本解析失败时带提示重试 1 次。
199        const MAX_RETRIES: usize = 1;
200        let mut current_prompt = prompt;
201
202        for attempt in 0..=MAX_RETRIES {
203            let result = chat_structured(
204                self.llm.as_ref(),
205                Some(queries_tool()),
206                vec![Message::human(&current_prompt)],
207            )
208            .await
209            .map_err(|e| MultiQueryError::LLMError(e.to_string()))?;
210
211            // 优先 tool_calls:查询字符串数组
212            if let Some(args) = &result.tool_args {
213                if let Some(queries) = parse_queries(args) {
214                    if !queries.is_empty() {
215                        return Ok(queries);
216                    }
217                }
218            }
219
220            // 文本兜底:逐行切,清理编号/引号/项目符号等脏文本
221            let queries = parse_query_lines(&result.content, self.config.num_queries);
222            if !queries.is_empty() {
223                return Ok(queries);
224            }
225
226            if attempt < MAX_RETRIES {
227                current_prompt = format!(
228                    "上次的输出不是有效的查询列表。请重新为原问题生成 {} 个不同的查询变体,\
229                     每行一个,不要编号、不要项目符号、不要解释或多余文字。\n\n原问题:{}\n\n\
230                     上次输出(无效):\n{}\n\n新的查询变体:",
231                    self.config.num_queries, original_query, result.content
232                );
233            }
234        }
235
236        Err(MultiQueryError::ParseError(
237            "LLM did not generate valid query variants".to_string(),
238        ))
239    }
240
241    /// 生成多个查询变体并分别检索，合并去重后返回最终文档
242    pub async fn retrieve_multi(&self, query: &str) -> Result<Vec<Document>, MultiQueryError> {
243        let queries = self.generate_queries(query).await?;
244
245        let all_queries: Vec<String> = std::iter::once(query.to_string()).chain(queries).collect();
246
247        let mut doc_scores: HashMap<String, (Document, f32)> = HashMap::new();
248
249        for q in &all_queries {
250            let results = self
251                .base_retriever
252                .retrieve_with_scores(q, self.config.k_per_query)
253                .await
254                .map_err(|e| MultiQueryError::RetrieverError(e.to_string()))?;
255
256            for result in results {
257                let doc_id = result
258                    .document
259                    .id
260                    .clone()
261                    .unwrap_or_else(|| doc_content_hash(&result.document.content));
262
263                doc_scores
264                    .entry(doc_id)
265                    .and_modify(|(_, score)| {
266                        // M3: use addition instead of no-op .max()
267                        *score += result.score;
268                    })
269                    .or_insert((result.document.clone(), result.score));
270            }
271        }
272
273        let mut scored_docs: Vec<(Document, f32)> = doc_scores
274            .values()
275            .map(|(doc, score)| (doc.clone(), *score))
276            .collect();
277
278        scored_docs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
279
280        let final_docs: Vec<Document> = scored_docs
281            .into_iter()
282            .take(self.config.final_k)
283            .map(|(doc, _)| doc)
284            .collect();
285
286        Ok(final_docs)
287    }
288
289    /// 生成多个查询变体并分别检索，合并去重后返回带分数的结果
290    pub async fn retrieve_multi_with_scores(
291        &self,
292        query: &str,
293    ) -> Result<Vec<SearchResult>, MultiQueryError> {
294        let queries = self.generate_queries(query).await?;
295
296        let all_queries: Vec<String> = std::iter::once(query.to_string()).chain(queries).collect();
297
298        let mut doc_scores: HashMap<String, (Document, f32, usize)> = HashMap::new();
299
300        for q in &all_queries {
301            let results = self
302                .base_retriever
303                .retrieve_with_scores(q, self.config.k_per_query)
304                .await
305                .map_err(|e| MultiQueryError::RetrieverError(e.to_string()))?;
306
307            for result in results {
308                let doc_id = result
309                    .document
310                    .id
311                    .clone()
312                    .unwrap_or_else(|| doc_content_hash(&result.document.content));
313
314                doc_scores
315                    .entry(doc_id)
316                    .and_modify(|(_, score, count)| {
317                        // M3: use addition instead of no-op .max()
318                        *score += result.score;
319                        *count += 1;
320                    })
321                    .or_insert((result.document.clone(), result.score, 1));
322            }
323        }
324
325        let mut scored_docs: Vec<SearchResult> = doc_scores
326            .values()
327            .map(|(doc, score, count)| {
328                let combined_score = score * (1.0 + 0.1 * *count as f32);
329                SearchResult {
330                    document: doc.clone(),
331                    score: combined_score,
332                }
333            })
334            .collect();
335
336        scored_docs.sort_by(|a, b| {
337            b.score
338                .partial_cmp(&a.score)
339                .unwrap_or(std::cmp::Ordering::Equal)
340        });
341
342        let final_results: Vec<SearchResult> =
343            scored_docs.into_iter().take(self.config.final_k).collect();
344
345        Ok(final_results)
346    }
347
348    /// 获取 LLM 生成的查询变体（不执行检索）
349    pub async fn get_generated_queries(&self, query: &str) -> Result<Vec<String>, MultiQueryError> {
350        self.generate_queries(query).await
351    }
352}
353
354/// 查询变体工具定义(P2-1):强制 LLM 输出查询字符串数组。
355fn queries_tool() -> ToolDefinition {
356    ToolDefinition::new(
357        "generate_queries",
358        "为原问题生成多个不同的检索查询变体,返回查询字符串数组",
359    )
360    .with_parameters(json!({
361        "type": "object",
362        "properties": {
363            "queries": {
364                "type": "array",
365                "items": { "type": "string" }
366            }
367        },
368        "required": ["queries"]
369    }))
370}
371
372/// 从 tool_call 参数中提取查询数组。
373fn parse_queries(args: &serde_json::Value) -> Option<Vec<String>> {
374    args.get("queries")?
375        .as_array()?
376        .iter()
377        .map(|v| v.as_str().map(|s| s.trim().to_string()))
378        .collect()
379}
380
381/// 文本行解析:清理编号"1. xxx"、项目符号、两端引号,避免脏文本当查询。
382///
383/// 只取前 `limit` 行(通常为 `num_queries`):LLM 常把查询排在最前,
384/// 末尾的解释性散文行会被截掉,不会混入查询列表。
385fn parse_query_lines(content: &str, limit: usize) -> Vec<String> {
386    content
387        .lines()
388        .map(|line| line.trim())
389        .filter(|line| !line.is_empty())
390        .map(|line| {
391            let stripped = line.trim_start_matches(['-', '•', '*', ' ']);
392            let stripped = stripped.trim_start_matches(|c: char| {
393                c.is_ascii_digit() || c == '.' || c == '、' || c == ')' || c == ' '
394            });
395            stripped
396                .trim_matches(['"', '\'', '“', '”'])
397                .trim()
398                .to_string()
399        })
400        .filter(|q| !q.is_empty())
401        .take(limit)
402        .collect()
403}
404
405/// 静态查询生成器（不依赖 LLM）
406#[allow(clippy::type_complexity)]
407pub struct StaticQueryGenerator {
408    expansions: Vec<Box<dyn Fn(&str) -> Vec<String> + Send + Sync>>,
409}
410
411impl StaticQueryGenerator {
412    /// 创建空的静态查询生成器
413    pub fn new() -> Self {
414        Self {
415            expansions: Vec::new(),
416        }
417    }
418
419    /// 添加同义词扩展规则
420    pub fn with_synonym_expansion(mut self, synonyms: HashMap<String, Vec<String>>) -> Self {
421        self.expansions.push(Box::new(move |query: &str| {
422            let mut expanded = Vec::new();
423            for (word, syns) in &synonyms {
424                if query.contains(word) {
425                    for syn in syns {
426                        expanded.push(query.replace(word, syn));
427                    }
428                }
429            }
430            expanded
431        }));
432        self
433    }
434
435    /// 添加前缀扩展规则
436    pub fn with_prefix_expansion(mut self, prefixes: Vec<String>) -> Self {
437        self.expansions.push(Box::new(move |query: &str| {
438            prefixes
439                .iter()
440                .map(|p| format!("{} {}", p, query))
441                .collect()
442        }));
443        self
444    }
445
446    /// 应用所有扩展规则生成查询变体
447    pub fn generate(&self, query: &str) -> Vec<String> {
448        self.expansions
449            .iter()
450            .flat_map(|exp| exp(query))
451            .filter(|q| q != query)
452            .collect()
453    }
454}
455
456impl Default for StaticQueryGenerator {
457    fn default() -> Self {
458        Self::new()
459    }
460}
461
462#[cfg(test)]
463mod tests {
464    use super::*;
465
466    #[test]
467    fn test_static_query_generator_synonym() {
468        let synonyms: HashMap<String, Vec<String>> = HashMap::from([(
469            "数据库".to_string(),
470            vec!["DB".to_string(), "存储".to_string()],
471        )]);
472
473        let generator = StaticQueryGenerator::new().with_synonym_expansion(synonyms);
474
475        let queries = generator.generate("数据库连接失败");
476
477        assert!(queries.contains(&"DB连接失败".to_string()));
478        assert!(queries.contains(&"存储连接失败".to_string()));
479    }
480
481    #[test]
482    fn test_static_query_generator_prefix() {
483        let generator = StaticQueryGenerator::new()
484            .with_prefix_expansion(vec!["如何".to_string(), "怎么".to_string()]);
485
486        let queries = generator.generate("处理错误");
487
488        assert!(queries.contains(&"如何 处理错误".to_string()));
489        assert!(queries.contains(&"怎么 处理错误".to_string()));
490    }
491
492    #[test]
493    fn test_multi_query_config() {
494        let config = MultiQueryConfig::new()
495            .with_num_queries(5)
496            .with_k_per_query(10)
497            .with_final_k(20);
498
499        assert_eq!(config.num_queries, 5);
500        assert_eq!(config.k_per_query, 10);
501        assert_eq!(config.final_k, 20);
502    }
503
504    #[test]
505    fn test_multi_query_config_default() {
506        let config = MultiQueryConfig::default();
507
508        assert_eq!(config.num_queries, 3);
509        assert_eq!(config.k_per_query, 5);
510        assert_eq!(config.final_k, 10);
511    }
512
513    /// P2-1: 查询工具定义携带 queries 数组 schema。
514    #[test]
515    fn test_queries_tool_schema() {
516        let tool = queries_tool();
517        assert_eq!(tool.function.name, "generate_queries");
518        let params = tool.function.parameters.expect("parameters should exist");
519        assert_eq!(params["properties"]["queries"]["type"], "array");
520    }
521
522    /// P2-1: tool_call 参数解析出查询数组。
523    #[test]
524    fn test_parse_queries() {
525        let args = json!({ "queries": ["数据库连接失败怎么办", "DB 连接错误排查"] });
526        let queries = parse_queries(&args).expect("should parse successfully");
527        assert_eq!(queries.len(), 2);
528        assert_eq!(queries[0], "数据库连接失败怎么办");
529    }
530
531    /// P2-1: 缺失 queries 键 → None。
532    #[test]
533    fn test_parse_queries_missing_key() {
534        let args = json!({ "other": 1 });
535        assert!(parse_queries(&args).is_none());
536    }
537
538    /// P2-1: 文本行解析清理编号/项目符号/引号,并按 limit 截断尾部散文。
539    #[test]
540    fn test_parse_query_lines_cleanup() {
541        let content = "1. 数据库连接失败\n- 如何排查 DB 错误\n• \"连接超时怎么办\"\n\n补充解释";
542        let queries = parse_query_lines(content, 3);
543        assert_eq!(
544            queries,
545            vec![
546                "数据库连接失败".to_string(),
547                "如何排查 DB 错误".to_string(),
548                "连接超时怎么办".to_string(),
549            ]
550        );
551    }
552
553    /// P2-1: 空文本 → 空数组。
554    #[test]
555    fn test_parse_query_lines_empty() {
556        assert!(parse_query_lines("  \n\n", 3).is_empty());
557    }
558
559    /// P2-1: 超过 limit 的额外行被截断。
560    #[test]
561    fn test_parse_query_lines_capped() {
562        let content = "a\nb\nc\nd";
563        assert_eq!(
564            parse_query_lines(content, 2),
565            vec!["a".to_string(), "b".to_string()]
566        );
567    }
568}
lc_rag/multi_query.rs

lc_rag/
multi_query.rs
