kv_cache.py

#!/usr/bin/env python3
"""KV-Cache System for MateCode - Core KV-Cache storage and retrieval

Based on Manus Context Engineering principles:
- Implement true KV-Cache storage (not just structural optimization)
- Cache entire prompts with same static prefix
- Provide basic cache hit rate statistics
- Simple TTL-based cache invalidation
"""

import hashlib
import json
import sqlite3
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict


@dataclass
class CacheStats:
    """Cache statistics"""
    hit_count: int = 0
    miss_count: int = 0
    total_queries: int = 0
    hit_rate: float = 0.0
    cache_size: int = 0  # number of entries


class KVCacheManager:
    """KV-Cache Manager - Core cache storage and retrieval

    Design principles:
    1. True KV-Cache storage (not just structural optimization)
    2. Simple cache key generation (static prefix hash + user ID)
    3. Basic TTL-based cache invalidation
    4. Minimal statistics tracking
    """

    def __init__(self, cache_dir: Optional[Path] = None):
        """Initialize KV-Cache manager

        Args:
            cache_dir: Directory for cache storage (default: ~/.matecode/kv_cache)
        """
        self.cache_dir = cache_dir or Path.home() / ".matecode" / "kv_cache"
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.db_path = self.cache_dir / "kv_cache.db"

        self._init_cache_db()
        self.stats = CacheStats()
        self._update_stats()

    def _init_cache_db(self):
        """Initialize cache database"""
        with sqlite3.connect(self.db_path) as conn:
            # Main cache table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS kv_cache (
                    cache_key TEXT PRIMARY KEY,
                    full_prompt TEXT NOT NULL,
                    static_prefix_hash TEXT NOT NULL,
                    user_id TEXT NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_accessed TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    access_count INTEGER DEFAULT 0,
                    ttl_seconds INTEGER DEFAULT 3600  -- 1 hour default TTL
                )
            """)

            # Index for efficient queries
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_cache_user_prefix
                ON kv_cache(user_id, static_prefix_hash)
            """)

            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_cache_last_accessed
                ON kv_cache(last_accessed)
            """)

            conn.commit()

    def generate_cache_key(self, static_prefix: str, user_id: str) -> str:
        """Generate cache key based on static prefix hash and user ID

        Args:
            static_prefix: The static prefix string (STATIC_SYSTEM_PREFIX)
            user_id: User/chat identifier

        Returns:
            Cache key string
        """
        # Hash the static prefix (MD5 is sufficient for this purpose)
        prefix_hash = hashlib.md5(static_prefix.encode()).hexdigest()[:16]

        # Combine with user ID for per-user caching
        return f"kv_cache:{user_id}:{prefix_hash}"

    def get_cached_prompt(self, cache_key: str) -> Optional[str]:
        """Get cached prompt by cache key

        Args:
            cache_key: Cache key generated by generate_cache_key()

        Returns:
            Cached prompt string if found and valid, None otherwise
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()

                # Check if cache entry exists and is not expired using SQLite datetime functions
                # This avoids timezone issues by doing all datetime math in SQLite
                cursor.execute("""
                    SELECT full_prompt, ttl_seconds, created_at,
                           (julianday('now') - julianday(created_at)) * 86400 as age_seconds
                    FROM kv_cache
                    WHERE cache_key = ?
                """, (cache_key,))

                row = cursor.fetchone()
                if not row:
                    self.stats.miss_count += 1
                    self.stats.total_queries += 1
                    return None

                ttl_seconds = row["ttl_seconds"]
                age_seconds = row["age_seconds"]

                if age_seconds > ttl_seconds:
                    # Cache expired, delete it
                    cursor.execute("DELETE FROM kv_cache WHERE cache_key = ?", (cache_key,))
                    conn.commit()
                    self.stats.miss_count += 1
                    self.stats.total_queries += 1
                    return None

                # Update access statistics
                cursor.execute("""
                    UPDATE kv_cache
                    SET last_accessed = CURRENT_TIMESTAMP,
                        access_count = access_count + 1
                    WHERE cache_key = ?
                """, (cache_key,))
                conn.commit()

                self.stats.hit_count += 1
                self.stats.total_queries += 1
                return row["full_prompt"]

        except Exception as e:
            print(f"Error retrieving from KV-Cache: {e}")
            self.stats.miss_count += 1
            self.stats.total_queries += 1
            return None

    def store_prompt(self, cache_key: str, full_prompt: str,
                     static_prefix: str, user_id: str,
                     ttl_seconds: int = 3600) -> bool:
        """Store prompt in cache

        Args:
            cache_key: Cache key generated by generate_cache_key()
            full_prompt: Complete prompt string to cache
            static_prefix: Static prefix string (for reference)
            user_id: User/chat identifier
            ttl_seconds: Time-to-live in seconds (default: 1 hour)

        Returns:
            True if successful, False otherwise
        """
        try:
            prefix_hash = hashlib.md5(static_prefix.encode()).hexdigest()[:16]

            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                # Insert or replace cache entry
                cursor.execute("""
                    INSERT OR REPLACE INTO kv_cache
                    (cache_key, full_prompt, static_prefix_hash, user_id, ttl_seconds)
                    VALUES (?, ?, ?, ?, ?)
                """, (cache_key, full_prompt, prefix_hash, user_id, ttl_seconds))

                conn.commit()

                # Update cache size stat
                self._update_stats()
                return True

        except Exception as e:
            print(f"Error storing to KV-Cache: {e}")
            return False

    def invalidate_cache(self, cache_key: Optional[str] = None,
                         user_id: Optional[str] = None) -> int:
        """Invalidate cache entries

        Args:
            cache_key: Specific cache key to invalidate (optional)
            user_id: Invalidate all cache entries for a user (optional)

        Returns:
            Number of cache entries invalidated
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                if cache_key:
                    cursor.execute("DELETE FROM kv_cache WHERE cache_key = ?", (cache_key,))
                elif user_id:
                    cursor.execute("DELETE FROM kv_cache WHERE user_id = ?", (user_id,))
                else:
                    # Invalidate all expired entries
                    cursor.execute("""
                        DELETE FROM kv_cache
                        WHERE (julianday('now') - julianday(created_at)) * 86400 > ttl_seconds
                    """)

                deleted_count = cursor.rowcount
                conn.commit()

                self._update_stats()
                return deleted_count

        except Exception as e:
            print(f"Error invalidating KV-Cache: {e}")
            return 0

    def _update_stats(self):
        """Update cache statistics"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT COUNT(*) FROM kv_cache")
                self.stats.cache_size = cursor.fetchone()[0]

                # Calculate hit rate
                total = self.stats.total_queries
                self.stats.hit_rate = self.stats.hit_count / total if total > 0 else 0.0

        except Exception as e:
            print(f"Error updating KV-Cache stats: {e}")

    def get_stats(self) -> Dict[str, Any]:
        """Get cache statistics

        Returns:
            Dictionary with cache statistics
        """
        self._update_stats()

        return {
            "hit_count": self.stats.hit_count,
            "miss_count": self.stats.miss_count,
            "total_queries": self.stats.total_queries,
            "hit_rate": self.stats.hit_rate,
            "cache_size": self.stats.cache_size,
            "cache_dir": str(self.cache_dir),
            "db_path": str(self.db_path),
        }

    def clear_cache(self) -> int:
        """Clear all cache entries

        Returns:
            Number of cache entries cleared
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT COUNT(*) FROM kv_cache")
                count = cursor.fetchone()[0]

                cursor.execute("DELETE FROM kv_cache")
                conn.commit()

                self._update_stats()
                return count

        except Exception as e:
            print(f"Error clearing KV-Cache: {e}")
            return 0


# Singleton instance pattern (similar to other memory modules)
_kv_cache_instance: Optional[KVCacheManager] = None

def get_kv_cache(cache_dir: Optional[Path] = None) -> KVCacheManager:
    """Get singleton KV-Cache instance

    Args:
        cache_dir: Directory for cache storage (optional)

    Returns:
        KV-Cache manager instance
    """
    global _kv_cache_instance
    if _kv_cache_instance is None:
        _kv_cache_instance = KVCacheManager(cache_dir)
    return _kv_cache_instance


# Convenience functions
def cache_prompt(static_prefix: str, user_id: str, full_prompt: str,
                 ttl_seconds: int = 3600) -> bool:
    """Convenience function: cache a prompt

    Args:
        static_prefix: Static prefix string
        user_id: User/chat identifier
        full_prompt: Complete prompt string
        ttl_seconds: Time-to-live in seconds

    Returns:
        True if successful
    """
    cache = get_kv_cache()
    cache_key = cache.generate_cache_key(static_prefix, user_id)
    return cache.store_prompt(cache_key, full_prompt, static_prefix, user_id, ttl_seconds)


def get_cached_prompt(static_prefix: str, user_id: str) -> Optional[str]:
    """Convenience function: get cached prompt

    Args:
        static_prefix: Static prefix string
        user_id: User/chat identifier

    Returns:
        Cached prompt string if found, None otherwise
    """
    cache = get_kv_cache()
    cache_key = cache.generate_cache_key(static_prefix, user_id)
    return cache.get_cached_prompt(cache_key)


def get_cache_stats() -> Dict[str, Any]:
    """Convenience function: get cache statistics

    Returns:
        Dictionary with cache statistics
    """
    return get_kv_cache().get_stats()