cache.py

"""Semantic caching system for API responses.

Supports exact match and fuzzy/semantic cache hits using embeddings.
"""

import sqlite3
import json
import time
import numpy as np
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
import hashlib
import threading
from abc import ABC, abstractmethod

from .utils import hash_content, normalize_prompt, format_messages, get_timestamp


@dataclass
class CacheEntry:
    """Represents a cached response."""
    key: str
    prompt_hash: str
    prompt_text: str
    response: str
    model: str
    embedding: Optional[List[float]]
    created_at: str
    expires_at: str
    hit_count: int = 0
    tokens_saved: int = 0
    cost_saved: float = 0.0


@dataclass
class CacheStats:
    """Cache statistics."""
    total_entries: int = 0
    total_hits: int = 0
    total_misses: int = 0
    exact_hits: int = 0
    semantic_hits: int = 0
    total_tokens_saved: int = 0
    total_cost_saved: float = 0.0

    @property
    def hit_rate(self) -> float:
        total = self.total_hits + self.total_misses
        return self.total_hits / total if total > 0 else 0.0


class EmbeddingProvider(ABC):
    """Abstract base class for embedding providers."""

    @abstractmethod
    def get_embedding(self, text: str) -> List[float]:
        """Get embedding vector for text."""
        pass


class SimpleEmbeddingProvider(EmbeddingProvider):
    """Simple TF-IDF-like embedding for offline use."""

    def __init__(self, dim: int = 256):
        self.dim = dim

    def get_embedding(self, text: str) -> List[float]:
        """Generate a simple hash-based embedding."""
        # Normalize text
        text = normalize_prompt(text)
        words = text.split()

        # Create embedding using word hashes
        embedding = [0.0] * self.dim
        for word in words:
            word_hash = int(hashlib.md5(word.encode()).hexdigest(), 16)
            for i in range(self.dim):
                embedding[i] += ((word_hash >> i) & 1) * 2 - 1

        # Normalize
        norm = np.linalg.norm(embedding)
        if norm > 0:
            embedding = [x / norm for x in embedding]

        return embedding


class OpenAIEmbeddingProvider(EmbeddingProvider):
    """OpenAI embedding provider."""

    def __init__(self, api_key: str, model: str = "text-embedding-3-small"):
        self.api_key = api_key
        self.model = model
        self._client = None

    @property
    def client(self):
        if self._client is None:
            try:
                from openai import OpenAI
                self._client = OpenAI(api_key=self.api_key)
            except ImportError:
                raise ImportError("openai package required for OpenAI embeddings")
        return self._client

    def get_embedding(self, text: str) -> List[float]:
        """Get embedding from OpenAI API."""
        response = self.client.embeddings.create(
            model=self.model,
            input=text[:8000]  # Truncate to avoid token limits
        )
        return response.data[0].embedding


class ResponseCache:
    """Semantic caching system for API responses."""

    def __init__(
        self,
        db_path: str = "cache.db",
        default_ttl: int = 3600,
        similarity_threshold: float = 0.92,
        embedding_provider: Optional[EmbeddingProvider] = None,
        max_entries: int = 10000
    ):
        """
        Initialize the cache.

        Args:
            db_path: Path to SQLite database
            default_ttl: Default time-to-live in seconds
            similarity_threshold: Minimum similarity for semantic match (0-1)
            embedding_provider: Provider for generating embeddings
            max_entries: Maximum cache entries before cleanup
        """
        self.db_path = Path(db_path)
        self.default_ttl = default_ttl
        self.similarity_threshold = similarity_threshold
        self.embedding_provider = embedding_provider or SimpleEmbeddingProvider()
        self.max_entries = max_entries
        self._lock = threading.Lock()
        self._stats = CacheStats()

        self._init_db()
        self._load_stats()

    def _init_db(self):
        """Initialize SQLite database."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS cache_entries (
                    key TEXT PRIMARY KEY,
                    prompt_hash TEXT NOT NULL,
                    prompt_text TEXT NOT NULL,
                    response TEXT NOT NULL,
                    model TEXT NOT NULL,
                    embedding BLOB,
                    created_at TEXT NOT NULL,
                    expires_at TEXT NOT NULL,
                    hit_count INTEGER DEFAULT 0,
                    tokens_saved INTEGER DEFAULT 0,
                    cost_saved REAL DEFAULT 0.0
                )
            """)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_prompt_hash 
                ON cache_entries(prompt_hash)
            """)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_expires_at 
                ON cache_entries(expires_at)
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS cache_stats (
                    id INTEGER PRIMARY KEY CHECK (id = 1),
                    total_hits INTEGER DEFAULT 0,
                    total_misses INTEGER DEFAULT 0,
                    exact_hits INTEGER DEFAULT 0,
                    semantic_hits INTEGER DEFAULT 0,
                    total_tokens_saved INTEGER DEFAULT 0,
                    total_cost_saved REAL DEFAULT 0.0
                )
            """)
            conn.execute("""
                INSERT OR IGNORE INTO cache_stats (id) VALUES (1)
            """)
            conn.commit()

    def _load_stats(self):
        """Load stats from database."""
        with sqlite3.connect(self.db_path) as conn:
            row = conn.execute("""
                SELECT total_hits, total_misses, exact_hits, semantic_hits,
                       total_tokens_saved, total_cost_saved
                FROM cache_stats WHERE id = 1
            """).fetchone()
            if row:
                self._stats.total_hits = row[0]
                self._stats.total_misses = row[1]
                self._stats.exact_hits = row[2]
                self._stats.semantic_hits = row[3]
                self._stats.total_tokens_saved = row[4]
                self._stats.total_cost_saved = row[5]

            count = conn.execute(
                "SELECT COUNT(*) FROM cache_entries"
            ).fetchone()[0]
            self._stats.total_entries = count

    def _save_stats(self):
        """Save stats to database."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                UPDATE cache_stats SET
                    total_hits = ?,
                    total_misses = ?,
                    exact_hits = ?,
                    semantic_hits = ?,
                    total_tokens_saved = ?,
                    total_cost_saved = ?
                WHERE id = 1
            """, (
                self._stats.total_hits,
                self._stats.total_misses,
                self._stats.exact_hits,
                self._stats.semantic_hits,
                self._stats.total_tokens_saved,
                self._stats.total_cost_saved
            ))
            conn.commit()

    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
        """Calculate cosine similarity between two vectors."""
        a_arr = np.array(a)
        b_arr = np.array(b)
        dot = np.dot(a_arr, b_arr)
        norm_a = np.linalg.norm(a_arr)
        norm_b = np.linalg.norm(b_arr)
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)

    def _is_expired(self, expires_at: str) -> bool:
        """Check if entry is expired."""
        exp_time = datetime.fromisoformat(expires_at)
        return datetime.utcnow() > exp_time

    def get(
        self,
        prompt: Union[str, List[Dict[str, str]]],
        model: str = "",
        semantic_search: bool = True
    ) -> Optional[Tuple[str, str]]:
        """
        Get cached response for prompt.

        Args:
            prompt: The prompt string or messages list
            model: Optional model filter
            semantic_search: Whether to use semantic matching

        Returns:
            Tuple of (response, match_type) or None if not found
        """
        with self._lock:
            # Normalize prompt
            if isinstance(prompt, list):
                prompt_text = format_messages(prompt)
            else:
                prompt_text = prompt

            prompt_hash = hash_content(normalize_prompt(prompt_text))

            with sqlite3.connect(self.db_path) as conn:
                # Try exact match first
                query = """
                    SELECT key, response, expires_at, embedding
                    FROM cache_entries
                    WHERE prompt_hash = ?
                """
                params = [prompt_hash]

                if model:
                    query += " AND model = ?"
                    params.append(model)

                rows = conn.execute(query, params).fetchall()

                for row in rows:
                    key, response, expires_at, _ = row
                    if not self._is_expired(expires_at):
                        self._record_hit(conn, key, "exact")
                        return (response, "exact")

                # Try semantic match if enabled
                if semantic_search:
                    prompt_embedding = self.embedding_provider.get_embedding(prompt_text)

                    # Get all non-expired entries
                    query = """
                        SELECT key, response, expires_at, embedding
                        FROM cache_entries
                        WHERE expires_at > ?
                    """
                    params = [datetime.utcnow().isoformat()]

                    if model:
                        query += " AND model = ?"
                        params.append(model)

                    rows = conn.execute(query, params).fetchall()

                    best_match = None
                    best_similarity = 0.0

                    for row in rows:
                        key, response, expires_at, embedding_blob = row
                        if embedding_blob:
                            stored_embedding = json.loads(embedding_blob)
                            similarity = self._cosine_similarity(
                                prompt_embedding, stored_embedding
                            )
                            if similarity > best_similarity:
                                best_similarity = similarity
                                best_match = (key, response)

                    if best_match and best_similarity >= self.similarity_threshold:
                        self._record_hit(conn, best_match[0], "semantic")
                        return (best_match[1], f"semantic:{best_similarity:.3f}")

            # Cache miss
            self._stats.total_misses += 1
            self._save_stats()
            return None

    def _record_hit(
        self,
        conn: sqlite3.Connection,
        key: str,
        hit_type: str
    ):
        """Record a cache hit."""
        conn.execute(
            "UPDATE cache_entries SET hit_count = hit_count + 1 WHERE key = ?",
            (key,)
        )
        conn.commit()

        self._stats.total_hits += 1
        if hit_type == "exact":
            self._stats.exact_hits += 1
        else:
            self._stats.semantic_hits += 1
        self._save_stats()

    def set(
        self,
        prompt: Union[str, List[Dict[str, str]]],
        response: str,
        model: str = "",
        ttl: Optional[int] = None,
        tokens_saved: int = 0,
        cost_saved: float = 0.0
    ) -> str:
        """
        Cache a response.

        Args:
            prompt: The prompt string or messages list
            response: The response to cache
            model: The model used
            ttl: Time-to-live in seconds (uses default if None)
            tokens_saved: Estimated tokens this cache will save
            cost_saved: Estimated cost this cache will save

        Returns:
            Cache entry key
        """
        with self._lock:
            # Normalize prompt
            if isinstance(prompt, list):
                prompt_text = format_messages(prompt)
            else:
                prompt_text = prompt

            prompt_hash = hash_content(normalize_prompt(prompt_text))
            key = hash_content(f"{prompt_hash}:{model}:{time.time()}")

            # Generate embedding
            embedding = self.embedding_provider.get_embedding(prompt_text)
            embedding_blob = json.dumps(embedding)

            # Calculate expiry
            ttl = ttl or self.default_ttl
            created_at = datetime.utcnow()
            expires_at = created_at + timedelta(seconds=ttl)

            with sqlite3.connect(self.db_path) as conn:
                conn.execute("""
                    INSERT OR REPLACE INTO cache_entries
                    (key, prompt_hash, prompt_text, response, model, embedding,
                     created_at, expires_at, hit_count, tokens_saved, cost_saved)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?)
                """, (
                    key, prompt_hash, prompt_text, response, model, embedding_blob,
                    created_at.isoformat(), expires_at.isoformat(),
                    tokens_saved, cost_saved
                ))
                conn.commit()

            self._stats.total_entries += 1

            # Cleanup if needed
            if self._stats.total_entries > self.max_entries:
                self._cleanup()

            return key

    def _cleanup(self):
        """Remove expired entries and oldest entries if over limit."""
        with sqlite3.connect(self.db_path) as conn:
            # Remove expired
            conn.execute(
                "DELETE FROM cache_entries WHERE expires_at < ?",
                (datetime.utcnow().isoformat(),)
            )

            # Remove oldest if still over limit
            count = conn.execute(
                "SELECT COUNT(*) FROM cache_entries"
            ).fetchone()[0]

            if count > self.max_entries:
                to_remove = count - self.max_entries + 100  # Remove extra buffer
                conn.execute("""
                    DELETE FROM cache_entries WHERE key IN (
                        SELECT key FROM cache_entries
                        ORDER BY hit_count ASC, created_at ASC
                        LIMIT ?
                    )
                """, (to_remove,))

            conn.commit()

            self._stats.total_entries = conn.execute(
                "SELECT COUNT(*) FROM cache_entries"
            ).fetchone()[0]

    def invalidate(self, key: str) -> bool:
        """Invalidate a specific cache entry."""
        with self._lock:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.execute(
                    "DELETE FROM cache_entries WHERE key = ?", (key,)
                )
                conn.commit()
                if cursor.rowcount > 0:
                    self._stats.total_entries -= 1
                    return True
                return False

    def clear(self):
        """Clear all cache entries."""
        with self._lock:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("DELETE FROM cache_entries")
                conn.commit()
            self._stats.total_entries = 0

    def get_stats(self) -> CacheStats:
        """Get cache statistics."""
        return self._stats

    def record_savings(self, tokens: int, cost: float):
        """Record tokens and cost saved from cache hit."""
        with self._lock:
            self._stats.total_tokens_saved += tokens
            self._stats.total_cost_saved += cost
            self._save_stats()