memory-search: add --query mode and prompt key boost
Two changes: 1. New -q/--query flag for direct search without hook machinery. Useful for debugging: memory-search -q inner-life-sexuality-intimacy shows seeds, spread results, and rankings. 2. Prompt key boost: when the current prompt contains a node key (>=5 chars) as a substring, boost that term by +10.0. This ensures explicit mentions fire as strong seeds for spread, while the graph still determines what gets pulled in. Co-Authored-By: ProofOfConcept <poc@bcachefs.org>
This commit is contained in:
parent
5024cf7002
commit
1da712874b
3 changed files with 100 additions and 8 deletions
|
|
@ -39,6 +39,10 @@ struct Args {
|
|||
#[arg(long, default_value = "5")]
|
||||
max_results: usize,
|
||||
|
||||
/// Search query (bypasses stashed input, uses this as the prompt)
|
||||
#[arg(long, short)]
|
||||
query: Option<String>,
|
||||
|
||||
/// Algorithm pipeline stages: e.g. spread spectral,k=20 spread,max_hops=4
|
||||
/// Default: spread.
|
||||
pipeline: Vec<String>,
|
||||
|
|
@ -61,6 +65,12 @@ fn main() {
|
|||
return;
|
||||
}
|
||||
|
||||
// --query mode: skip all hook/context machinery, just search
|
||||
if let Some(ref query_str) = args.query {
|
||||
run_query_mode(query_str, &args);
|
||||
return;
|
||||
}
|
||||
|
||||
let input = if args.hook {
|
||||
// Hook mode: read from stdin, stash for later debug runs
|
||||
let mut buf = String::new();
|
||||
|
|
@ -206,6 +216,22 @@ fn main() {
|
|||
}
|
||||
}
|
||||
|
||||
// Boost node keys that appear as substrings in the current prompt.
|
||||
// Makes explicit mentions strong seeds for spread — the graph
|
||||
// determines what gets pulled in, this just ensures the seed fires.
|
||||
{
|
||||
let prompt_lower = prompt.to_lowercase();
|
||||
for (key, node) in &store.nodes {
|
||||
if node.deleted { continue; }
|
||||
let key_lower = key.to_lowercase();
|
||||
if key_lower.len() < 5 { continue; }
|
||||
if prompt_lower.contains(&key_lower) {
|
||||
*terms.entry(key_lower).or_insert(0.0) += 10.0;
|
||||
if debug { println!("[memory-search] prompt key boost: {} (+10.0)", key); }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if debug {
|
||||
println!("[memory-search] {} terms total", terms.len());
|
||||
let mut by_weight: Vec<_> = terms.iter().collect();
|
||||
|
|
@ -336,6 +362,65 @@ fn main() {
|
|||
cleanup_stale_files(&state_dir, Duration::from_secs(86400));
|
||||
}
|
||||
|
||||
/// Direct query mode: search for a term without hook/stash machinery.
|
||||
fn run_query_mode(query: &str, args: &Args) {
|
||||
let store = match poc_memory::store::Store::load() {
|
||||
Ok(s) => s,
|
||||
Err(e) => { eprintln!("failed to load store: {}", e); return; }
|
||||
};
|
||||
|
||||
// Build terms from the query string
|
||||
let mut terms: BTreeMap<String, f64> = BTreeMap::new();
|
||||
let prompt_terms = search::extract_query_terms(query, 8);
|
||||
for word in prompt_terms.split_whitespace() {
|
||||
terms.entry(word.to_lowercase()).or_insert(1.0);
|
||||
}
|
||||
|
||||
// Also check for exact node key match (the query itself, lowercased)
|
||||
let query_lower = query.to_lowercase();
|
||||
for (key, node) in &store.nodes {
|
||||
if node.deleted { continue; }
|
||||
if key.to_lowercase() == query_lower {
|
||||
terms.insert(query_lower.clone(), 10.0);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
println!("[query] terms: {:?}", terms);
|
||||
|
||||
if terms.is_empty() {
|
||||
println!("[query] no terms extracted");
|
||||
return;
|
||||
}
|
||||
|
||||
let graph = poc_memory::graph::build_graph_fast(&store);
|
||||
let (seeds, direct_hits) = search::match_seeds(&terms, &store);
|
||||
|
||||
println!("[query] {} seeds", seeds.len());
|
||||
let mut sorted = seeds.clone();
|
||||
sorted.sort_by(|a, b| b.1.total_cmp(&a.1));
|
||||
for (key, score) in sorted.iter().take(20) {
|
||||
let marker = if direct_hits.contains(key) { "→" } else { " " };
|
||||
println!(" {} {:.4} {}", marker, score, key);
|
||||
}
|
||||
|
||||
let pipeline: Vec<AlgoStage> = if args.pipeline.is_empty() {
|
||||
vec![AlgoStage::parse("spread").unwrap()]
|
||||
} else {
|
||||
args.pipeline.iter()
|
||||
.filter_map(|a| AlgoStage::parse(a).ok())
|
||||
.collect()
|
||||
};
|
||||
|
||||
let max_results = args.max_results.max(25);
|
||||
let results = search::run_pipeline(&pipeline, seeds, &graph, &store, true, max_results);
|
||||
|
||||
println!("\n[query] top {} results:", results.len().min(25));
|
||||
for (i, (key, score)) in results.iter().take(25).enumerate() {
|
||||
let marker = if direct_hits.contains(key) { "→" } else { " " };
|
||||
println!(" {:2}. {} [{:.4}] {}", i + 1, marker, score, key);
|
||||
}
|
||||
}
|
||||
|
||||
/// Split context output into chunks of approximately `max_bytes`, breaking
|
||||
/// at section boundaries ("--- KEY (group) ---" lines).
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue