query_expansion/
query_expansion.rs1use 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 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 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 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 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}