#!/usr/bin/env python3 """query_memory.py — Query memory system for adaptive retrieval. The system remembers what kinds of queries it gets, what retrieval strategy worked, and adapts its gate logic accordingly. This prevents dead weight — memories get tested based on usage patterns, not just random exploration. Research basis: - IGMiRAG (arXiv:2602.07525): Intuition-guided retrieval - Task-Adaptive Retrieval (arXiv:2604.07863): Learning from interaction history - MemCon (arXiv:2607.13591): Learned policy for consolidation/forgetting - Agent-Native Memory (arXiv:2606.24775): Workload-adaptive routing Run under ~/.venvs/memory/bin/python3. """ import json import os import sys from collections import Counter, defaultdict from datetime import datetime, timezone from typing import Dict, List, Optional, Tuple QUERY_MEMORY_DIR = os.path.expanduser("~/.config/opencode/memory/query_memory") QUERY_HISTORY_FILE = os.path.join(QUERY_MEMORY_DIR, "history.jsonl") TOPIC_STATS_FILE = os.path.join(QUERY_MEMORY_DIR, "topic_stats.json") # Topic keywords for classification TOPIC_KEYWORDS = { "guitar": ["guitar", "music", "teaching", "course", "lesson", "chord", "scale"], "coding": ["code", "python", "script", "function", "debug", "build", "deploy"], "memory": ["memory", "recall", "remember", "forget", "consolidat", "neuron"], "persona": ["persona", "identity", "voice", "constitution", "character"], "project": ["project", "task", "goal", "deadline", "milestone", "progress"], "ai": ["ai", "model", "llm", "agent", "neural", "training", "inference"], "personal": ["nick", "nottingham", "coffee", "coconut", "dog", "ttrpg"], "system": ["system", "config", "install", "setup", "troubleshoot", "error"], "research": ["research", "paper", "study", "analysis", "survey", "evidence"], "creative": ["write", "design", "create", "build", "make", "art", "music"], } class QueryMemory: """Query memory system for adaptive retrieval.""" def __init__(self): os.makedirs(QUERY_MEMORY_DIR, exist_ok=True) self.history: List[dict] = self._load_history() self.topic_stats: Dict[str, dict] = self._load_topic_stats() def _load_history(self) -> List[dict]: """Load query history from disk.""" if not os.path.exists(QUERY_HISTORY_FILE): return [] entries = [] with open(QUERY_HISTORY_FILE) as f: for line in f: try: entries.append(json.loads(line.strip())) except json.JSONDecodeError: continue return entries def _save_history(self, entry: dict): """Append a query entry to history.""" with open(QUERY_HISTORY_FILE, "a") as f: f.write(json.dumps(entry) + "\n") def _load_topic_stats(self) -> Dict[str, dict]: """Load topic statistics from disk.""" if os.path.exists(TOPIC_STATS_FILE): try: with open(TOPIC_STATS_FILE) as f: return json.load(f) except json.JSONDecodeError: return {} return {} def _save_topic_stats(self): """Save topic statistics to disk.""" with open(TOPIC_STATS_FILE, "w") as f: json.dump(self.topic_stats, f, indent=2) def classify_query(self, query: str) -> str: """Classify a query into a topic. Returns the topic with the most keyword matches. Falls back to "general" if no strong match. """ query_lower = query.lower() scores = {} for topic, keywords in TOPIC_KEYWORDS.items(): score = sum(1 for kw in keywords if kw in query_lower) if score > 0: scores[topic] = score if not scores: return "general" return max(scores, key=scores.get) def record_query(self, query: str, strategy: str, success: bool): """Record a query and its outcome. Args: query: The query text strategy: "consolidated" or "full_retrieval" success: Whether the retrieval produced useful results """ topic = self.classify_query(query) entry = { "query": query[:100], "topic": topic, "strategy": strategy, "success": success, "timestamp": datetime.now(timezone.utc).isoformat(), } self.history.append(entry) self._save_history(entry) # Update topic stats if topic not in self.topic_stats: self.topic_stats[topic] = { "total_queries": 0, "successful_queries": 0, "consolidated_hits": 0, "full_retrieval_hits": 0, "exploration_rate": 0.30, # default } stats = self.topic_stats[topic] stats["total_queries"] += 1 if success: stats["successful_queries"] += 1 if strategy == "consolidated": stats["consolidated_hits"] += 1 else: stats["full_retrieval_hits"] += 1 # Adaptive exploration rate: decrease for well-tested topics if stats["total_queries"] >= 10: success_rate = stats["successful_queries"] / stats["total_queries"] if success_rate > 0.8: stats["exploration_rate"] = max(0.10, stats["exploration_rate"] - 0.05) elif success_rate < 0.5: stats["exploration_rate"] = min(0.50, stats["exploration_rate"] + 0.05) self._save_topic_stats() def get_topic_stats(self) -> Dict[str, dict]: """Get topic statistics.""" return self.topic_stats def adaptive_exploration_rate(self, query: str) -> float: """Get the exploration rate for a query based on its topic. Well-tested topics have lower exploration rates. Untested topics have higher exploration rates. """ topic = self.classify_query(query) if topic in self.topic_stats: return self.topic_stats[topic].get("exploration_rate", 0.30) return 0.50 # untested topics get high exploration def get_dead_weight_topics(self) -> List[str]: """Get topics with low proven uses (dead weight).""" dead = [] for topic, stats in self.topic_stats.items(): if stats["total_queries"] < 5: dead.append(topic) return dead def get_health_report(self) -> dict: """Get a health report of the query memory system.""" total_queries = len(self.history) topics = len(self.topic_stats) dead_topics = len(self.get_dead_weight_topics()) return { "total_queries": total_queries, "topics_tracked": topics, "dead_weight_topics": dead_topics, "topic_stats": self.topic_stats, } # ═══════════════════════════════════════════════════════════════════════════════ # CLI # ═══════════════════════════════════════════════════════════════════════════════ def main(): import argparse parser = argparse.ArgumentParser(description="Query memory system") sub = parser.add_subparsers(dest="cmd") # record --strategy --success p_rec = sub.add_parser("record", help="Record a query") p_rec.add_argument("query", nargs="+") p_rec.add_argument("--strategy", default="consolidated") p_rec.add_argument("--success", action="store_true") # stats sub.add_parser("stats", help="Show topic statistics") # rate p_rate = sub.add_parser("rate", help="Get exploration rate for a query") p_rate.add_argument("query", nargs="+") # health sub.add_parser("health", help="Show health report") args = parser.parse_args() if not args.cmd: parser.print_help() sys.exit(1) qm = QueryMemory() if args.cmd == "record": query = " ".join(args.query) qm.record_query(query, args.strategy, args.success) print(f"Recorded: {query[:50]} → {args.strategy} ({'success' if args.success else 'fail'})") elif args.cmd == "stats": stats = qm.get_topic_stats() if stats: print("Topic Statistics:") for topic, s in stats.items(): print(f" {topic}: {s['total_queries']} queries, {s['successful_queries']} success, rate={s['exploration_rate']:.2f}") else: print("No topic statistics yet.") elif args.cmd == "rate": query = " ".join(args.query) rate = qm.adaptive_exploration_rate(query) topic = qm.classify_query(query) print(f"Query: {query[:50]}") print(f"Topic: {topic}") print(f"Exploration rate: {rate:.2f}") elif args.cmd == "health": report = qm.get_health_report() print("Query Memory Health:") print(f" Total queries: {report['total_queries']}") print(f" Topics tracked: {report['topics_tracked']}") print(f" Dead weight topics: {report['dead_weight_topics']}") if __name__ == "__main__": main()