Skip to main content

query_expansion/
query_expansion.rs

1//! Query expansion for improved search coverage - fully self-contained.
2//!
3//! Run: `cargo run --example query_expansion`
4
5use anyhow::Result;
6use qmd::{GenerationEngine, QueryType, Queryable, Store, expand_query_simple};
7
8const SAMPLE_DOCS: &[(&str, &str)] = &[
9    ("rust-basics.md", include_str!("data/rust-basics.md")),
10    ("error-handling.md", include_str!("data/error-handling.md")),
11];
12
13fn main() -> Result<()> {
14    let db_path = std::env::temp_dir().join("qmd_expansion.db");
15    let _ = std::fs::remove_file(&db_path);
16    let store = Store::open(&db_path)?;
17
18    let now = chrono::Utc::now().to_rfc3339();
19    for (name, content) in SAMPLE_DOCS {
20        let hash = Store::hash_content(content);
21        let title = Store::extract_title(content);
22        store.insert_content(&hash, content, &now)?;
23        store.insert_document("samples", name, &title, &hash, &now, &now)?;
24    }
25
26    let query = "rust error handling";
27    println!("Query: '{}'\n", query);
28
29    // Simple expansion
30    println!("Simple expansion:");
31    for q in expand_query_simple(query) {
32        let t = match q.query_type {
33            QueryType::Lex => "LEX",
34            QueryType::Vec => "VEC",
35            QueryType::Hyde => "HYD",
36        };
37        println!("  [{}] {}", t, q.text);
38    }
39
40    // Search with expanded queries
41    println!("\nSearch results:");
42    for q in expand_query_simple(query) {
43        if q.query_type == QueryType::Lex {
44            let n = store.search_fts(&q.text, 5, None)?.len();
45            println!("  '{}': {} results", q.text, n);
46        }
47    }
48
49    // LLM expansion
50    println!("\nLLM expansion:");
51    if GenerationEngine::is_available() {
52        let engine = GenerationEngine::load_default()?;
53        for q in engine.expand_query(query, true)? {
54            println!("  [{:?}] {}", q.query_type, q.text);
55        }
56    } else {
57        println!("  (not available)");
58    }
59
60    // Manual construction
61    println!("\nManual:");
62    for q in [
63        Queryable::lex("rust error"),
64        Queryable::vec("exception handling"),
65    ] {
66        println!("  [{:?}] {}", q.query_type, q.text);
67    }
68
69    let _ = std::fs::remove_file(&db_path);
70    Ok(())
71}