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:
Kent Overstreet 2026-03-13 15:26:35 -04:00
parent 5024cf7002
commit 1da712874b
3 changed files with 100 additions and 8 deletions

View file

@ -9,8 +9,8 @@
// - Conversation fragment selection (for observation agent) // - Conversation fragment selection (for observation agent)
use crate::graph::Graph; use crate::graph::Graph;
use super::llm;
use crate::spectral; use crate::spectral;
use super::llm;
use crate::store::{self, Store, new_relation, RelationType}; use crate::store::{self, Store, new_relation, RelationType};
use regex::Regex; use regex::Regex;
@ -940,14 +940,21 @@ fn run_cycle(
depth_db.save(&mut store); depth_db.save(&mut store);
// Recompute spectral if anything changed // Recompute spectral at most once per hour — O(n³) is expensive at 14k+ nodes
if total_applied > 0 { if total_applied > 0 {
eprintln!("\n Recomputing spectral embedding..."); let stale = spectral::embedding_path()
.metadata()
.and_then(|m| m.modified())
.map(|t| t.elapsed().unwrap_or_default() > std::time::Duration::from_secs(3600))
.unwrap_or(true);
if stale {
eprintln!("\n Recomputing spectral embedding (>1h stale)...");
let graph = store.build_graph(); let graph = store.build_graph();
let result = spectral::decompose(&graph, 8); let result = spectral::decompose(&graph, 8);
let emb = spectral::to_embedding(&result); let emb = spectral::to_embedding(&result);
spectral::save_embedding(&emb).ok(); spectral::save_embedding(&emb).ok();
} }
}
let graph = store.build_graph(); let graph = store.build_graph();
let metrics_after = GraphMetrics::from_graph(&store, &graph); let metrics_after = GraphMetrics::from_graph(&store, &graph);

View file

@ -39,6 +39,10 @@ struct Args {
#[arg(long, default_value = "5")] #[arg(long, default_value = "5")]
max_results: usize, 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 /// Algorithm pipeline stages: e.g. spread spectral,k=20 spread,max_hops=4
/// Default: spread. /// Default: spread.
pipeline: Vec<String>, pipeline: Vec<String>,
@ -61,6 +65,12 @@ fn main() {
return; 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 { let input = if args.hook {
// Hook mode: read from stdin, stash for later debug runs // Hook mode: read from stdin, stash for later debug runs
let mut buf = String::new(); 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 { if debug {
println!("[memory-search] {} terms total", terms.len()); println!("[memory-search] {} terms total", terms.len());
let mut by_weight: Vec<_> = terms.iter().collect(); let mut by_weight: Vec<_> = terms.iter().collect();
@ -336,6 +362,65 @@ fn main() {
cleanup_stale_files(&state_dir, Duration::from_secs(86400)); 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 /// Split context output into chunks of approximately `max_bytes`, breaking
/// at section boundaries ("--- KEY (group) ---" lines). /// at section boundaries ("--- KEY (group) ---" lines).

View file

@ -42,7 +42,7 @@ pub struct SpectralEmbedding {
pub coords: HashMap<String, Vec<f64>>, pub coords: HashMap<String, Vec<f64>>,
} }
fn embedding_path() -> PathBuf { pub fn embedding_path() -> PathBuf {
crate::store::memory_dir().join("spectral-embedding.json") crate::store::memory_dir().join("spectral-embedding.json")
} }