lc_rag/
multi_query.rs

1// src/retrieval/multi_query.rs
2//! MultiQueryRetriever implementation
3//!
4//! Uses an LLM to generate multiple query variants, improving retrieval recall.
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: Replaces `DefaultHasher` with FNV-1a 64-bit (the std internal algorithm is not
22/// guaranteed stable across processes; FNV-1a is a fully-specified deterministic hash).
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 error type
31#[derive(Debug)]
32#[non_exhaustive]
33pub enum MultiQueryError {
34    /// LLM error
35    LLMError(String),
36
37    /// Retriever error
38    RetrieverError(String),
39
40    /// Parse error
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 configuration
57pub struct MultiQueryConfig {
58    /// Number of generated queries
59    pub num_queries: usize,
60
61    /// Number of documents returned per query
62    pub k_per_query: usize,
63
64    /// Number of documents returned in the final result
65    pub final_k: usize,
66
67    /// The query-generation 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    /// Creates a `MultiQueryConfig` with default configuration
84    pub fn new() -> Self {
85        Self::default()
86    }
87
88    /// Sets the number of generated queries
89    pub fn with_num_queries(mut self, n: usize) -> Self {
90        self.num_queries = n;
91        self
92    }
93
94    /// Sets the number of documents returned per query
95    pub fn with_k_per_query(mut self, k: usize) -> Self {
96        self.k_per_query = k;
97        self
98    }
99
100    /// Sets the number of documents returned in the final result
101    pub fn with_final_k(mut self, k: usize) -> Self {
102        self.final_k = k;
103        self
104    }
105
106    /// Sets the query-generation 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/// Uses an LLM to generate multiple query variants, retrieves with the base retriever for
126/// each one, then merges and dedups the results before returning.
127pub struct MultiQueryRetriever {
128    /// The LLM used to generate query variants
129    ///
130    /// P0-3: no longer hardcodes `OpenAIChat`; accepts any LLM implementing `BaseChatModel`.
131    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
132
133    /// The base retriever
134    base_retriever: Arc<dyn RetrieverTrait>,
135
136    /// Configuration
137    config: MultiQueryConfig,
138}
139
140impl MultiQueryRetriever {
141    /// Creates a MultiQueryRetriever (accepting any LLM implementing `BaseChatModel`)
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: Builds from an already-wrapped `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    /// Sets the MultiQuery configuration
167    pub fn with_config(mut self, config: MultiQueryConfig) -> Self {
168        self.config = config;
169        self
170    }
171
172    /// Sets the number of generated queries
173    pub fn with_num_queries(mut self, n: usize) -> Self {
174        self.config.num_queries = n;
175        self
176    }
177
178    /// Sets the number of documents returned per query
179    pub fn with_k_per_query(mut self, k: usize) -> Self {
180        self.config.k_per_query = k;
181        self
182    }
183
184    /// Sets the number of documents returned in the final result
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: Prefer the structured query list from tool_calls; on text-parse failure,
199        // retry once with a hint.
200        const MAX_RETRIES: usize = 1;
201        let mut current_prompt = prompt;
202
203        for attempt in 0..=MAX_RETRIES {
204            let result = chat_structured(
205                self.llm.as_ref(),
206                Some(queries_tool()),
207                vec![Message::human(&current_prompt)],
208            )
209            .await
210            .map_err(|e| MultiQueryError::LLMError(e.to_string()))?;
211
212            // Prefer tool_calls: the query string array
213            if let Some(args) = &result.tool_args {
214                if let Some(queries) = parse_queries(args) {
215                    if !queries.is_empty() {
216                        return Ok(queries);
217                    }
218                }
219            }
220
221            // Text fallback: split line by line, cleaning numbered/quote/bullet noise
222            let queries = parse_query_lines(&result.content, self.config.num_queries);
223            if !queries.is_empty() {
224                return Ok(queries);
225            }
226
227            if attempt < MAX_RETRIES {
228                current_prompt = format!(
229                    "上次的输出不是有效的查询列表。请重新为原问题生成 {} 个不同的查询变体,\
230                     每行一个,不要编号、不要项目符号、不要解释或多余文字。\n\n原问题:{}\n\n\
231                     上次输出(无效):\n{}\n\n新的查询变体:",
232                    self.config.num_queries, original_query, result.content
233                );
234            }
235        }
236
237        Err(MultiQueryError::ParseError(
238            "LLM did not generate valid query variants".to_string(),
239        ))
240    }
241
242    /// Generates multiple query variants, retrieves each one, and returns the merged deduped documents
243    pub async fn retrieve_multi(&self, query: &str) -> Result<Vec<Document>, MultiQueryError> {
244        let queries = self.generate_queries(query).await?;
245
246        let all_queries: Vec<String> = std::iter::once(query.to_string()).chain(queries).collect();
247
248        let mut doc_scores: HashMap<String, (Document, f32)> = HashMap::new();
249
250        for q in &all_queries {
251            let results = self
252                .base_retriever
253                .retrieve_with_scores(q, self.config.k_per_query)
254                .await
255                .map_err(|e| MultiQueryError::RetrieverError(e.to_string()))?;
256
257            for result in results {
258                let doc_id = result
259                    .document
260                    .id
261                    .clone()
262                    .unwrap_or_else(|| doc_content_hash(&result.document.content));
263
264                doc_scores
265                    .entry(doc_id)
266                    .and_modify(|(_, score)| {
267                        // M3: use addition instead of no-op .max()
268                        *score += result.score;
269                    })
270                    .or_insert((result.document.clone(), result.score));
271            }
272        }
273
274        let mut scored_docs: Vec<(Document, f32)> = doc_scores
275            .values()
276            .map(|(doc, score)| (doc.clone(), *score))
277            .collect();
278
279        scored_docs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
280
281        let final_docs: Vec<Document> = scored_docs
282            .into_iter()
283            .take(self.config.final_k)
284            .map(|(doc, _)| doc)
285            .collect();
286
287        Ok(final_docs)
288    }
289
290    /// Generates multiple query variants, retrieves each one, and returns the merged deduped results with scores
291    pub async fn retrieve_multi_with_scores(
292        &self,
293        query: &str,
294    ) -> Result<Vec<SearchResult>, MultiQueryError> {
295        let queries = self.generate_queries(query).await?;
296
297        let all_queries: Vec<String> = std::iter::once(query.to_string()).chain(queries).collect();
298
299        let mut doc_scores: HashMap<String, (Document, f32, usize)> = HashMap::new();
300
301        for q in &all_queries {
302            let results = self
303                .base_retriever
304                .retrieve_with_scores(q, self.config.k_per_query)
305                .await
306                .map_err(|e| MultiQueryError::RetrieverError(e.to_string()))?;
307
308            for result in results {
309                let doc_id = result
310                    .document
311                    .id
312                    .clone()
313                    .unwrap_or_else(|| doc_content_hash(&result.document.content));
314
315                doc_scores
316                    .entry(doc_id)
317                    .and_modify(|(_, score, count)| {
318                        // M3: use addition instead of no-op .max()
319                        *score += result.score;
320                        *count += 1;
321                    })
322                    .or_insert((result.document.clone(), result.score, 1));
323            }
324        }
325
326        let mut scored_docs: Vec<SearchResult> = doc_scores
327            .values()
328            .map(|(doc, score, count)| {
329                let combined_score = score * (1.0 + 0.1 * *count as f32);
330                SearchResult {
331                    document: doc.clone(),
332                    score: combined_score,
333                }
334            })
335            .collect();
336
337        scored_docs.sort_by(|a, b| {
338            b.score
339                .partial_cmp(&a.score)
340                .unwrap_or(std::cmp::Ordering::Equal)
341        });
342
343        let final_results: Vec<SearchResult> =
344            scored_docs.into_iter().take(self.config.final_k).collect();
345
346        Ok(final_results)
347    }
348
349    /// Returns the LLM-generated query variants (without retrieving)
350    pub async fn get_generated_queries(&self, query: &str) -> Result<Vec<String>, MultiQueryError> {
351        self.generate_queries(query).await
352    }
353}
354
355/// Query-variant tool definition (P2-1): forces the LLM to output a query string array.
356fn queries_tool() -> ToolDefinition {
357    ToolDefinition::new(
358        "generate_queries",
359        "为原问题生成多个不同的检索查询变体,返回查询字符串数组",
360    )
361    .with_parameters(json!({
362        "type": "object",
363        "properties": {
364            "queries": {
365                "type": "array",
366                "items": { "type": "string" }
367            }
368        },
369        "required": ["queries"]
370    }))
371}
372
373/// Extracts the query array from tool_call arguments.
374fn parse_queries(args: &serde_json::Value) -> Option<Vec<String>> {
375    args.get("queries")?
376        .as_array()?
377        .iter()
378        .map(|v| v.as_str().map(|s| s.trim().to_string()))
379        .collect()
380}
381
382/// Text-line parsing: strips numbering ("1. xxx"), bullets, and surrounding quotes so dirty
383/// text never becomes a query.
384///
385/// Takes only the first `limit` lines (usually `num_queries`): the LLM typically puts the
386/// queries first, and trailing explanatory prose gets truncated instead of leaking in.
387fn parse_query_lines(content: &str, limit: usize) -> Vec<String> {
388    content
389        .lines()
390        .map(|line| line.trim())
391        .filter(|line| !line.is_empty())
392        .map(|line| {
393            let stripped = line.trim_start_matches(['-', '•', '*', ' ']);
394            let stripped = stripped.trim_start_matches(|c: char| {
395                c.is_ascii_digit() || c == '.' || c == '、' || c == ')' || c == ' '
396            });
397            stripped
398                .trim_matches(['"', '\'', '“', '”'])
399                .trim()
400                .to_string()
401        })
402        .filter(|q| !q.is_empty())
403        .take(limit)
404        .collect()
405}
406
407/// A static query generator (no LLM dependency)
408#[allow(clippy::type_complexity)]
409pub struct StaticQueryGenerator {
410    expansions: Vec<Box<dyn Fn(&str) -> Vec<String> + Send + Sync>>,
411}
412
413impl StaticQueryGenerator {
414    /// Creates an empty static query generator
415    pub fn new() -> Self {
416        Self {
417            expansions: Vec::new(),
418        }
419    }
420
421    /// Adds a synonym-expansion rule
422    pub fn with_synonym_expansion(mut self, synonyms: HashMap<String, Vec<String>>) -> Self {
423        self.expansions.push(Box::new(move |query: &str| {
424            let mut expanded = Vec::new();
425            for (word, syns) in &synonyms {
426                if query.contains(word) {
427                    for syn in syns {
428                        expanded.push(query.replace(word, syn));
429                    }
430                }
431            }
432            expanded
433        }));
434        self
435    }
436
437    /// Adds a prefix-expansion rule
438    pub fn with_prefix_expansion(mut self, prefixes: Vec<String>) -> Self {
439        self.expansions.push(Box::new(move |query: &str| {
440            prefixes
441                .iter()
442                .map(|p| format!("{} {}", p, query))
443                .collect()
444        }));
445        self
446    }
447
448    /// Applies all expansion rules to generate query variants
449    pub fn generate(&self, query: &str) -> Vec<String> {
450        self.expansions
451            .iter()
452            .flat_map(|exp| exp(query))
453            .filter(|q| q != query)
454            .collect()
455    }
456}
457
458impl Default for StaticQueryGenerator {
459    fn default() -> Self {
460        Self::new()
461    }
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467
468    #[test]
469    fn test_static_query_generator_synonym() {
470        let synonyms: HashMap<String, Vec<String>> = HashMap::from([(
471            "数据库".to_string(),
472            vec!["DB".to_string(), "存储".to_string()],
473        )]);
474
475        let generator = StaticQueryGenerator::new().with_synonym_expansion(synonyms);
476
477        let queries = generator.generate("数据库连接失败");
478
479        assert!(queries.contains(&"DB连接失败".to_string()));
480        assert!(queries.contains(&"存储连接失败".to_string()));
481    }
482
483    #[test]
484    fn test_static_query_generator_prefix() {
485        let generator = StaticQueryGenerator::new()
486            .with_prefix_expansion(vec!["如何".to_string(), "怎么".to_string()]);
487
488        let queries = generator.generate("处理错误");
489
490        assert!(queries.contains(&"如何 处理错误".to_string()));
491        assert!(queries.contains(&"怎么 处理错误".to_string()));
492    }
493
494    #[test]
495    fn test_multi_query_config() {
496        let config = MultiQueryConfig::new()
497            .with_num_queries(5)
498            .with_k_per_query(10)
499            .with_final_k(20);
500
501        assert_eq!(config.num_queries, 5);
502        assert_eq!(config.k_per_query, 10);
503        assert_eq!(config.final_k, 20);
504    }
505
506    #[test]
507    fn test_multi_query_config_default() {
508        let config = MultiQueryConfig::default();
509
510        assert_eq!(config.num_queries, 3);
511        assert_eq!(config.k_per_query, 5);
512        assert_eq!(config.final_k, 10);
513    }
514
515    /// P2-1: The query tool definition carries a queries-array schema.
516    #[test]
517    fn test_queries_tool_schema() {
518        let tool = queries_tool();
519        assert_eq!(tool.function.name, "generate_queries");
520        let params = tool.function.parameters.expect("parameters should exist");
521        assert_eq!(params["properties"]["queries"]["type"], "array");
522    }
523
524    /// P2-1: The tool_call arguments parse into a query array.
525    #[test]
526    fn test_parse_queries() {
527        let args = json!({ "queries": ["数据库连接失败怎么办", "DB 连接错误排查"] });
528        let queries = parse_queries(&args).expect("should parse successfully");
529        assert_eq!(queries.len(), 2);
530        assert_eq!(queries[0], "数据库连接失败怎么办");
531    }
532
533    /// P2-1: A missing `queries` key → None.
534    #[test]
535    fn test_parse_queries_missing_key() {
536        let args = json!({ "other": 1 });
537        assert!(parse_queries(&args).is_none());
538    }
539
540    /// P2-1: Text-line parsing strips numbering/bullets/quotes and caps trailing prose by limit.
541    #[test]
542    fn test_parse_query_lines_cleanup() {
543        let content = "1. 数据库连接失败\n- 如何排查 DB 错误\n• \"连接超时怎么办\"\n\n补充解释";
544        let queries = parse_query_lines(content, 3);
545        assert_eq!(
546            queries,
547            vec![
548                "数据库连接失败".to_string(),
549                "如何排查 DB 错误".to_string(),
550                "连接超时怎么办".to_string(),
551            ]
552        );
553    }
554
555    /// P2-1: Empty text → empty array.
556    #[test]
557    fn test_parse_query_lines_empty() {
558        assert!(parse_query_lines("  \n\n", 3).is_empty());
559    }
560
561    /// P2-1: Extra lines beyond the limit are truncated.
562    #[test]
563    fn test_parse_query_lines_capped() {
564        let content = "a\nb\nc\nd";
565        assert_eq!(
566            parse_query_lines(content, 2),
567            vec!["a".to_string(), "b".to_string()]
568        );
569    }
570}
lc_rag/multi_query.rs

lc_rag/
multi_query.rs
