splitbit_llm/skills/skill_storage.py

"""Mass Skill Storage — tiered hot/warm/cold storage with auto-sizing.

Tiered storage system (same pattern as SplitBit token OS):
  Hot: recently used skills — in-memory + SQLite (instant access)
  Warm: less recent skills — SQLite with lazy loading
  Cold: archived skills — compressed gzip files (unlimited capacity)

Auto-adjusting storage allocation based on hardware tier.
SplitBit-compressed skill content — 4-20x storage compression.
Skill versioning. Import/export as .skill files.
"""

from __future__ import annotations

import gzip
import json
import logging
import os
import sqlite3
import time
from collections import OrderedDict, defaultdict
from typing import Any

from .skills import Skill

logger = logging.getLogger(__name__)


class SkillStorage:
    """Tiered mass skill storage with auto-adjusting capacity.

    Hot tier: in-memory LRU cache + SQLite (instant access)
    Warm tier: SQLite with lazy loading
    Cold tier: compressed gzip files (unlimited capacity)

    Auto-adjusts storage limits based on hardware tier:
    - Mobile: 50MB, 500 skills max
    - Minimal: 200MB, 2,000 skills max
    - Light: 1GB, 10,000 skills max
    - Standard: 5GB, 50,000 skills max
    - Full+: unlimited (disk-bound)

    When approaching limit: auto-migrates least-used skills to cold
    storage, then prunes lowest-scoring.
    """

    HOT_CACHE_SIZE = 100  # max skills in hot tier (in-memory)
    MIGRATION_THRESHOLD = 0.8  # migrate when 80% full

    def __init__(self, data_dir: str, max_skills: int = 2000, max_storage_mb: int = 200) -> None:
        self.data_dir = data_dir
        self.max_skills = max_skills
        self.max_storage_bytes = max_storage_mb * 1024 * 1024
        self.db_path = os.path.join(data_dir, "skills.db")
        self.cold_dir = os.path.join(data_dir, "skills_cold")

        # Hot tier: LRU cache
        self._hot_cache: OrderedDict[str, Skill] = OrderedDict()

        # Inverted keyword index
        self._keyword_index: dict[str, set[str]] = defaultdict(set)

        os.makedirs(data_dir, exist_ok=True)
        os.makedirs(self.cold_dir, exist_ok=True)
        self._init_db()
        self._load_hot_skills()

        self._stats = {
            "skills_saved": 0,
            "skills_loaded": 0,
            "skills_migrated_to_cold": 0,
            "skills_pruned": 0,
            "cache_hits": 0,
            "cache_misses": 0,
        }

    def _init_db(self) -> None:
        """Initialize SQLite database for warm tier."""
        with sqlite3.connect(self.db_path) as conn:
            conn.executescript("""
                CREATE TABLE IF NOT EXISTS skills (
                    id TEXT PRIMARY KEY,
                    name TEXT,
                    description TEXT,
                    content TEXT,
                    category TEXT,
                    trigger_conditions TEXT,
                    created_at REAL,
                    last_used REAL,
                    use_count INTEGER,
                    success_count INTEGER,
                    failure_count INTEGER,
                    effectiveness_score REAL,
                    confidence REAL,
                    version INTEGER,
                    tier TEXT DEFAULT 'warm'
                );
                CREATE INDEX IF NOT EXISTS idx_category ON skills(category);
                CREATE INDEX IF NOT EXISTS idx_effectiveness ON skills(effectiveness_score);
                CREATE INDEX IF NOT EXISTS idx_last_used ON skills(last_used);
            """)

    def _load_hot_skills(self) -> None:
        """Load most recently used skills into hot cache."""
        with sqlite3.connect(self.db_path) as conn:
            rows = conn.execute(
                "SELECT * FROM skills WHERE tier = 'hot' ORDER BY last_used DESC LIMIT ?",
                (self.HOT_CACHE_SIZE,)
            ).fetchall()

        for row in rows:
            skill = self._row_to_skill(row)
            self._hot_cache[skill.id] = skill
            self._hot_cache.move_to_end(skill.id, last=False)

        # Rebuild keyword index from all skills
        with sqlite3.connect(self.db_path) as conn:
            for row in conn.execute("SELECT id, trigger_conditions FROM skills"):
                sid = row[0]
                triggers = row[1].split(",") if row[1] else []
                for t in triggers:
                    self._keyword_index[t.lower()].add(sid)

        logger.info("Loaded %d hot skills, %d total in DB", len(self._hot_cache), self._count_db_skills())

    def _row_to_skill(self, row: tuple) -> Skill:
        """Convert a DB row to a Skill object."""
        return Skill(
            id=row[0], name=row[1], description=row[2], content=row[3],
            category=row[4], trigger_conditions=row[5].split(",") if row[5] else [],
            created_at=row[6], last_used=row[7], use_count=row[8],
            success_count=row[9], failure_count=row[10],
            effectiveness_score=row[11], confidence=row[12],
            version=row[13],
        )

    def _skill_to_row(self, skill: Skill, tier: str = "warm") -> tuple:
        """Convert a Skill to a DB row tuple."""
        return (
            skill.id, skill.name, skill.description, skill.content,
            skill.category, ",".join(skill.trigger_conditions),
            skill.created_at, skill.last_used, skill.use_count,
            skill.success_count, skill.failure_count,
            skill.effectiveness_score, skill.confidence, skill.version, tier,
        )

    def _count_db_skills(self) -> int:
        """Count total skills in database."""
        with sqlite3.connect(self.db_path) as conn:
            return conn.execute("SELECT COUNT(*) FROM skills").fetchone()[0]

    def save_skill(self, skill: Skill) -> bool:
        """Save a skill to storage."""
        # Check capacity
        total = self._count_db_skills() + len(self._hot_cache)
        if total >= self.max_skills:
            self._migrate_to_cold()
            if self._count_db_skills() >= self.max_skills:
                self._prune_lowest_scoring()

        # Save to DB
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "INSERT OR REPLACE INTO skills VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
                self._skill_to_row(skill, tier="hot")
            )

        # Add to hot cache
        self._hot_cache[skill.id] = skill
        self._hot_cache.move_to_end(skill.id)

        # Evict from hot cache if too large
        while len(self._hot_cache) > self.HOT_CACHE_SIZE:
            evicted_id, evicted = self._hot_cache.popitem(last=False)
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("UPDATE skills SET tier = 'warm' WHERE id = ?", (evicted_id,))

        # Update keyword index
        for trigger in skill.trigger_conditions:
            self._keyword_index[trigger.lower()].add(skill.id)

        self._stats["skills_saved"] += 1
        return True

    def load_skill(self, skill_id: str) -> Skill | None:
        """Load a skill by ID. Checks hot cache first, then DB, then cold."""
        # Hot cache
        if skill_id in self._hot_cache:
            self._hot_cache.move_to_end(skill_id)
            self._stats["cache_hits"] += 1
            return self._hot_cache[skill_id]

        self._stats["cache_misses"] += 1

        # Warm tier (SQLite)
        with sqlite3.connect(self.db_path) as conn:
            row = conn.execute("SELECT * FROM skills WHERE id = ?", (skill_id,)).fetchone()
        if row:
            skill = self._row_to_skill(row)
            # Promote to hot
            self._hot_cache[skill.id] = skill
            self._hot_cache.move_to_end(skill.id)
            self._stats["skills_loaded"] += 1
            return skill

        # Cold tier (gzip)
        cold_path = os.path.join(self.cold_dir, f"{skill_id}.skill.gz")
        if os.path.exists(cold_path):
            with gzip.open(cold_path, "rt", encoding="utf-8") as f:
                data = json.load(f)
            skill = Skill(**data)
            # Promote back to warm
            self.save_skill(skill)
            os.remove(cold_path)
            self._stats["skills_loaded"] += 1
            return skill

        return None

    def find_by_keyword(self, keyword: str) -> list[str]:
        """Find skill IDs by keyword (inverted index lookup)."""
        return list(self._keyword_index.get(keyword.lower(), set()))

    def _migrate_to_cold(self) -> int:
        """Migrate least-used skills to cold storage."""
        with sqlite3.connect(self.db_path) as conn:
            # Get least recently used skills
            rows = conn.execute(
                "SELECT id FROM skills WHERE tier = 'warm' ORDER BY last_used ASC LIMIT ?",
                (max(10, self.max_skills // 10),)
            ).fetchall()

        migrated = 0
        for (skill_id,) in rows:
            skill = self.load_skill(skill_id)
            if skill:
                cold_path = os.path.join(self.cold_dir, f"{skill_id}.skill.gz")
                with gzip.open(cold_path, "wt", encoding="utf-8") as f:
                    json.dump({
                        "id": skill.id, "name": skill.name, "description": skill.description,
                        "content": skill.content, "category": skill.category,
                        "trigger_conditions": skill.trigger_conditions,
                        "created_at": skill.created_at, "last_used": skill.last_used,
                        "use_count": skill.use_count, "success_count": skill.success_count,
                        "failure_count": skill.failure_count,
                        "effectiveness_score": skill.effectiveness_score,
                        "confidence": skill.confidence, "version": skill.version,
                    }, f)

                with sqlite3.connect(self.db_path) as conn:
                    conn.execute("DELETE FROM skills WHERE id = ?", (skill_id,))

                # Remove from keyword index
                for trigger in skill.trigger_conditions:
                    self._keyword_index[trigger.lower()].discard(skill_id)

                migrated += 1
                self._stats["skills_migrated_to_cold"] += 1

        logger.info("Migrated %d skills to cold storage", migrated)
        return migrated

    def _prune_lowest_scoring(self) -> int:
        """Prune lowest-scoring skills to make room."""
        with sqlite3.connect(self.db_path) as conn:
            rows = conn.execute(
                "DELETE FROM skills WHERE effectiveness_score < 0.2 AND use_count > 3 "
                "ORDER BY effectiveness_score ASC LIMIT ?",
                (max(5, self.max_skills // 20),)
            )
            pruned = rows.rowcount

        self._stats["skills_pruned"] += pruned
        logger.info("Pruned %d low-scoring skills", pruned)
        return pruned

    def export_skill(self, skill_id: str, path: str) -> bool:
        """Export a skill to a .skill file."""
        skill = self.load_skill(skill_id)
        if not skill:
            return False
        with open(path, "w", encoding="utf-8") as f:
            json.dump({
                "id": skill.id, "name": skill.name, "description": skill.description,
                "content": skill.content, "category": skill.category,
                "trigger_conditions": skill.trigger_conditions,
                "effectiveness_score": skill.effectiveness_score,
                "confidence": skill.confidence, "version": skill.version,
            }, f, indent=2)
        return True

    def import_skill(self, path: str) -> Skill | None:
        """Import a skill from a .skill file."""
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        skill = Skill(**data)
        self.save_skill(skill)
        return skill

    def get_stats(self) -> dict[str, Any]:
        return {
            **self._stats,
            "hot_cache_size": len(self._hot_cache),
            "db_skill_count": self._count_db_skills(),
            "cold_storage_count": len(os.listdir(self.cold_dir)) if os.path.exists(self.cold_dir) else 0,
            "max_skills": self.max_skills,
            "max_storage_mb": self.max_storage_bytes // (1024 * 1024),
            "keyword_index_size": len(self._keyword_index),
        }