consciousness/src/bin/memory-search.rs

205 lines
6.4 KiB
Rust
Raw Normal View History

// memory-search: combined hook for session context loading + ambient memory retrieval
//
// On first prompt per session: loads full memory context (identity, journal, etc.)
// On subsequent prompts: searches memory for relevant entries
// On post-compaction: reloads full context
//
// Reads JSON from stdin (Claude Code UserPromptSubmit hook format),
// outputs results for injection into the conversation.
use std::collections::HashSet;
use std::fs;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::{Duration, SystemTime};
fn main() {
let mut input = String::new();
io::stdin().read_to_string(&mut input).unwrap_or_default();
let json: serde_json::Value = match serde_json::from_str(&input) {
Ok(v) => v,
Err(_) => return,
};
let prompt = json["prompt"].as_str().unwrap_or("");
let session_id = json["session_id"].as_str().unwrap_or("");
if prompt.is_empty() || session_id.is_empty() {
return;
}
let state_dir = PathBuf::from("/tmp/claude-memory-search");
fs::create_dir_all(&state_dir).ok();
// Detect post-compaction reload
let is_compaction = prompt.contains("continued from a previous conversation");
// First prompt or post-compaction: load full context
let cookie_path = state_dir.join(format!("cookie-{}", session_id));
let is_first = !cookie_path.exists();
if is_first || is_compaction {
// Create/touch the cookie
let cookie = if is_first {
let c = generate_cookie();
fs::write(&cookie_path, &c).ok();
c
} else {
fs::read_to_string(&cookie_path).unwrap_or_default().trim().to_string()
};
// Load full memory context
if let Ok(output) = Command::new("poc-memory").args(["load-context"]).output() {
if output.status.success() {
let ctx = String::from_utf8_lossy(&output.stdout);
if !ctx.trim().is_empty() {
print!("{}", ctx);
}
}
}
// On first prompt, also bump lookup counter for the cookie
let _ = cookie; // used for tagging below
}
// Always do ambient search (skip on very short or system prompts)
let word_count = prompt.split_whitespace().count();
if word_count < 3 {
return;
}
for prefix in &["is AFK", "You're on your own", "IRC mention"] {
if prompt.starts_with(prefix) {
return;
}
}
let query = extract_query_terms(prompt, 3);
if query.is_empty() {
return;
}
let output = Command::new("poc-memory")
.args(["search", &query])
.output();
let search_output = match output {
Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout).to_string(),
_ => return,
};
if search_output.trim().is_empty() {
return;
}
let cookie = fs::read_to_string(&cookie_path).unwrap_or_default().trim().to_string();
let seen = load_seen(&state_dir, session_id);
let mut result_output = String::new();
let mut count = 0;
let max_entries = 5;
for line in search_output.lines() {
if count >= max_entries { break; }
let trimmed = line.trim();
if trimmed.is_empty() { continue; }
if let Some(key) = extract_key_from_line(trimmed) {
if seen.contains(&key) { continue; }
mark_seen(&state_dir, session_id, &key);
result_output.push_str(line);
result_output.push('\n');
count += 1;
} else if count > 0 {
result_output.push_str(line);
result_output.push('\n');
}
}
if count == 0 { return; }
println!("Recalled memories [{}]:", cookie);
print!("{}", result_output);
// Clean up stale state files (opportunistic)
cleanup_stale_files(&state_dir, Duration::from_secs(86400));
}
fn extract_query_terms(text: &str, max_terms: usize) -> String {
const STOP_WORDS: &[&str] = &[
"the", "a", "an", "is", "are", "was", "were", "do", "does", "did",
"have", "has", "had", "will", "would", "could", "should", "can",
"may", "might", "shall", "been", "being", "to", "of", "in", "for",
"on", "with", "at", "by", "from", "as", "but", "or", "and", "not",
"no", "if", "then", "than", "that", "this", "it", "its", "my",
"your", "our", "we", "you", "i", "me", "he", "she", "they", "them",
"what", "how", "why", "when", "where", "about", "just", "let",
"want", "tell", "show", "think", "know", "see", "look", "make",
"get", "go", "some", "any", "all", "very", "really", "also", "too",
"so", "up", "out", "here", "there",
];
text.to_lowercase()
.split(|c: char| !c.is_alphanumeric())
.filter(|w| !w.is_empty() && w.len() > 2 && !STOP_WORDS.contains(w))
.take(max_terms)
.collect::<Vec<_>>()
.join(" ")
}
fn extract_key_from_line(line: &str) -> Option<String> {
let after_bracket = line.find("] ")?;
let rest = &line[after_bracket + 2..];
let key_end = rest.find(" (c").unwrap_or(rest.len());
let key = rest[..key_end].trim();
if key.is_empty() || !key.contains('.') {
None
} else {
Some(key.to_string())
}
}
fn generate_cookie() -> String {
uuid::Uuid::new_v4().as_simple().to_string()[..12].to_string()
}
fn load_seen(dir: &Path, session_id: &str) -> HashSet<String> {
let path = dir.join(format!("seen-{}", session_id));
if path.exists() {
fs::read_to_string(path)
.unwrap_or_default()
.lines()
.map(|s| s.to_string())
.collect()
} else {
HashSet::new()
}
}
fn mark_seen(dir: &Path, session_id: &str, key: &str) {
let path = dir.join(format!("seen-{}", session_id));
if let Ok(mut f) = fs::OpenOptions::new().create(true).append(true).open(path) {
writeln!(f, "{}", key).ok();
}
}
fn cleanup_stale_files(dir: &Path, max_age: Duration) {
let entries = match fs::read_dir(dir) {
Ok(e) => e,
Err(_) => return,
};
let cutoff = SystemTime::now() - max_age;
for entry in entries.flatten() {
if let Ok(meta) = entry.metadata() {
if let Ok(modified) = meta.modified() {
if modified < cutoff {
fs::remove_file(entry.path()).ok();
}
}
}
}
}