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

lc_rag/
multi_query.rs
