rune/shared/shard_manager.py

#!/usr/bin/env python3
"""shard_manager.py — Shard-based storage for unbounded memory growth.

Splits storage by topic into separate SQLite shard files.
No ceiling: new topics = new shards = new files = new disk space.

Each shard has its own FTS5 + sqlite-vec index for fast retrieval.
Query routing sends queries to the correct shard based on topic.
Results merge across shards when needed.

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

import json
import os
import sqlite3
import sys
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple

SHARD_DIR = os.path.expanduser("~/.config/opencode/memory/shards")
SHARD_INDEX = os.path.expanduser("~/.config/opencode/memory/shard_index.json")


class ShardManager:
    """Manages topic-based SQLite shards for unbounded memory growth."""

    def __init__(self):
        os.makedirs(SHARD_DIR, exist_ok=True)
        self.index = self._load_index()

    def _load_index(self) -> Dict[str, str]:
        """Load shard index from disk."""
        if os.path.exists(SHARD_INDEX):
            try:
                with open(SHARD_INDEX) as f:
                    return json.load(f)
            except json.JSONDecodeError:
                return {}
        return {}

    def _save_index(self):
        """Save shard index to disk."""
        with open(SHARD_INDEX, "w") as f:
            json.dump(self.index, f, indent=2)

    def create_shard(self, topic: str) -> str:
        """Create a new shard for a topic. Returns shard path."""
        shard_name = f"shard_{topic.replace(' ', '_').replace('/', '_')}.db"
        shard_path = os.path.join(SHARD_DIR, shard_name)

        if os.path.exists(shard_path):
            return shard_path

        db = sqlite3.connect(shard_path)
        db.executescript("""
            CREATE TABLE IF NOT EXISTS entries (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                text TEXT NOT NULL,
                importance REAL DEFAULT 0.5,
                created_at TEXT NOT NULL,
                last_accessed TEXT,
                topic TEXT DEFAULT '',
                source TEXT DEFAULT '',
                status TEXT DEFAULT 'active',
                wash_count INTEGER DEFAULT 0,
                last_washed TEXT,
                rock_status TEXT DEFAULT 'washable',
                tier TEXT DEFAULT 'warm'
            );
            CREATE VIRTUAL TABLE IF NOT EXISTS entries_fts USING fts5(
                text, content='', tokenize='porter unicode61'
            );
            CREATE INDEX IF NOT EXISTS idx_entries_status ON entries(status);
            CREATE INDEX IF NOT EXISTS idx_entries_topic ON entries(topic);
            CREATE INDEX IF NOT EXISTS idx_entries_wash_count ON entries(wash_count);
            CREATE INDEX IF NOT EXISTS idx_entries_tier ON entries(tier);
        """)
        db.close()

        self.index[topic] = shard_path
        self._save_index()
        return shard_path

    def route_query(self, query: str) -> str:
        """Route a query to the appropriate shard topic.

        Returns the topic name. Falls back to 'general' if no match.
        """
        query_lower = query.lower()

        # Check existing shards for keyword matches
        best_topic = None
        best_score = 0

        for topic in self.index.keys():
            topic_words = set(topic.lower().split("_"))
            query_words = set(query_lower.split())
            overlap = len(topic_words & query_words)
            if overlap > best_score:
                best_score = overlap
                best_topic = topic

        if best_topic and best_score > 0:
            return best_topic

        # Extract topic from query
        words = [w for w in query.split() if len(w) > 3][:3]
        if words:
            return "_".join(words).lower()

        return "general"

    def get_shard(self, topic: str) -> sqlite3.Connection:
        """Get or create a shard for a topic. Returns database connection."""
        if topic not in self.index:
            self.create_shard(topic)

        shard_path = self.index[topic]
        db = sqlite3.connect(shard_path)
        db.row_factory = sqlite3.Row
        return db

    def merge_results(self, results_list: List[List[dict]]) -> List[dict]:
        """Merge results from multiple shards, deduplicate, sort by relevance."""
        seen = set()
        merged = []

        for results in results_list:
            for r in results:
                key = r.get("text", "")[:50]  # deduplicate by text prefix
                if key not in seen:
                    seen.add(key)
                    merged.append(r)

        # Sort by relevance/wash_count
        merged.sort(key=lambda x: x.get("relevance", x.get("wash_count", 0)), reverse=True)
        return merged

    def get_shard_stats(self) -> Dict[str, int]:
        """Get entry counts per shard."""
        stats = {}
        for topic, path in self.index.items():
            try:
                db = sqlite3.connect(path)
                count = db.execute("SELECT COUNT(*) FROM entries WHERE status='active'").fetchone()[0]
                stats[topic] = count
                db.close()
            except Exception:
                stats[topic] = 0
        return stats

    def total_entries(self) -> int:
        """Get total entries across all shards."""
        stats = self.get_shard_stats()
        return sum(stats.values())


# ═══════════════════════════════════════════════════════════════════════════════
# CLI
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Shard manager for unbounded memory")
    sub = parser.add_subparsers(dest="cmd")

    sub.add_parser("stats", help="Show shard statistics")

    p_route = sub.add_parser("route", help="Route a query to a shard")
    p_route.add_argument("query", nargs="+")

    p_create = sub.add_parser("create", help="Create a shard for a topic")
    p_create.add_argument("topic")

    args = parser.parse_args()

    if not args.cmd:
        parser.print_help()
        sys.exit(1)

    sm = ShardManager()

    if args.cmd == "stats":
        stats = sm.get_shard_stats()
        print(f"Shard Statistics ({len(stats)} shards, {sm.total_entries()} total entries):")
        for topic, count in sorted(stats.items(), key=lambda x: x[1], reverse=True):
            print(f"  {topic}: {count} entries")

    elif args.cmd == "route":
        query = " ".join(args.query)
        topic = sm.route_query(query)
        print(f"Query: {query}")
        print(f"Route to: {topic}")

    elif args.cmd == "create":
        path = sm.create_shard(args.topic)
        print(f"Created shard: {path}")


if __name__ == "__main__":
    main()