rune/shared/query_memory.py

#!/usr/bin/env python3
"""query_memory.py — Query memory system for adaptive retrieval.

AUTHORITY-SWITCHED (Phase G strangler slice 1, doctrine step 8): the
classification core now lives in Rust (`src/rust/query_memory/src/lib.rs`,
cdylib `libquery_memory.so`). This module is a thin ctypes-backed shim over
`qm_loader.py` serving the FULL historical Python class API — the deployed
in-process consumers (`memory_store.py`, `growth_simulator.py`) call
`QueryMemory()` exactly as before and get the Rust core.

Parity evidence: FFI-gap parity EXACT-EQUAL on `get_topic_stats` /
`get_dead_weight_topics`; verification wall green (FFI smoke 41/41, parity
harness ALL GREEN, Python canary 12/12, Rust canary 27/27); shadow comparison
46/46 exact-equal vs the live production store.

The Python implementation this replaces is preserved in git history on branch
`phase-g-strangler-retrieval-index` (pre-shim commit); per deletion discipline
it is not kept as a silent fallback — a missing cdylib fails loudly.

Store resolution is unchanged: `$HOME/.config/opencode/memory/query_memory`
(the Rust core resolves `$HOME` natively; history.jsonl / topic_stats.json
formats are byte-parity-proven).

Run under ~/.venvs/memory/bin/python3.
"""

import json
import os
import sys
from datetime import datetime, timezone

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 (classification table — tie-break law: insertion order is
# part of the contract; the Rust core carries the identical table).
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"],
}

# Where the Rust core (qm_loader.py + target/release/libquery_memory.so) lives.
_QM_CORE_DIR = os.environ.get(
    "QM_CORE_DIR", "/mnt/workdrive/rune/src/rust/query_memory"
)

if _QM_CORE_DIR not in sys.path:
    sys.path.insert(0, _QM_CORE_DIR)

try:
    from qm_loader import QueryMemory as _RustQueryMemory
except ImportError as _e:  # pragma: no cover - deployment integrity failure
    raise ImportError(
        f"query_memory shim: cannot load the Rust core from {_QM_CORE_DIR!r} "
        f"({_e}). The cdylib must be built: `cargo build --release` in "
        f"{_QM_CORE_DIR}. No Python fallback exists by design."
    ) from _e


class QueryMemory(_RustQueryMemory):
    """Query memory system for adaptive retrieval (Rust core, Python API).

    API (identical to the historical Python implementation):
      * classify_query(query) -> str
      * adaptive_exploration_rate(query) -> float
      * record_query(query, strategy, success) -> None
      * get_topic_stats() -> dict (topic -> stats, insertion order)
      * get_dead_weight_topics() -> list[str]
      * get_health_report() -> dict
    """

    def __init__(self):
        # Match the historical constructor: ensure the store dir exists,
        # resolve the store from $HOME (qm_new's native behaviour).
        os.makedirs(QUERY_MEMORY_DIR, exist_ok=True)
        super().__init__()
        # Historical in-memory attribute: query history loaded at construction,
        # appended by record_query. (canary check 2 reads len(qm.history).)
        self.history = self._load_history()

    def _load_history(self):
        """Load query history from disk (historical shape: list of dicts)."""
        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 record_query(self, query: str, strategy: str, success: bool):
        """Record through the Rust core, then mirror the historical in-memory append."""
        super().record_query(query, strategy, success)
        self.history.append({
            "query": query[:100],
            "topic": self.classify_query(query),
            "strategy": strategy,
            "success": success,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })


# ═══════════════════════════════════════════════════════════════════════════════
# 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()