rune/shared/query_memory.py

#!/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 <query> --strategy <s> --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 <query>
    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()