ai-lib-python 0.5.0__py3-none-any.whl

ai_lib_python/cache/backends.py

@@ -0,0 +1,380 @@
+"""
+Cache backend implementations.
+
+Provides memory, disk, and null cache backends.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class CacheEntry:
+    """A cache entry with metadata.
+
+    Attributes:
+        value: Cached value
+        created_at: Creation timestamp
+        ttl: Time-to-live in seconds
+        hits: Number of cache hits
+    """
+
+    value: Any
+    created_at: float
+    ttl: float | None = None
+    hits: int = 0
+
+    @property
+    def is_expired(self) -> bool:
+        """Check if entry is expired."""
+        if self.ttl is None:
+            return False
+        return time.time() > self.created_at + self.ttl
+
+    @property
+    def age_seconds(self) -> float:
+        """Get age in seconds."""
+        return time.time() - self.created_at
+
+
+class CacheBackend(ABC):
+    """Abstract base class for cache backends."""
+
+    @abstractmethod
+    async def get(self, key: str) -> Any | None:
+        """Get a value from the cache.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            Cached value or None
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def set(
+        self, key: str, value: Any, ttl: float | None = None
+    ) -> None:
+        """Set a value in the cache.
+
+        Args:
+            key: Cache key
+            value: Value to cache
+            ttl: Time-to-live in seconds
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete(self, key: str) -> bool:
+        """Delete a value from the cache.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            True if deleted, False if not found
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def clear(self) -> None:
+        """Clear all cache entries."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def exists(self, key: str) -> bool:
+        """Check if a key exists in the cache.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            True if exists, False otherwise
+        """
+        raise NotImplementedError
+
+    async def close(self) -> None:
+        """Close the backend (cleanup)."""
+        pass
+
+
+class MemoryCache(CacheBackend):
+    """In-memory cache backend with TTL support.
+
+    Example:
+        >>> cache = MemoryCache(max_size=1000, default_ttl=3600)
+        >>> await cache.set("key", {"data": "value"})
+        >>> value = await cache.get("key")
+    """
+
+    def __init__(
+        self,
+        max_size: int = 1000,
+        default_ttl: float | None = None,
+    ) -> None:
+        """Initialize memory cache.
+
+        Args:
+            max_size: Maximum number of entries
+            default_ttl: Default TTL in seconds
+        """
+        self._cache: dict[str, CacheEntry] = {}
+        self._max_size = max_size
+        self._default_ttl = default_ttl
+        self._lock = asyncio.Lock()
+
+    async def get(self, key: str) -> Any | None:
+        """Get a value from the cache."""
+        async with self._lock:
+            entry = self._cache.get(key)
+            if entry is None:
+                return None
+
+            if entry.is_expired:
+                del self._cache[key]
+                return None
+
+            entry.hits += 1
+            return entry.value
+
+    async def set(
+        self, key: str, value: Any, ttl: float | None = None
+    ) -> None:
+        """Set a value in the cache."""
+        async with self._lock:
+            # Evict if at capacity
+            if len(self._cache) >= self._max_size and key not in self._cache:
+                self._evict_one()
+
+            self._cache[key] = CacheEntry(
+                value=value,
+                created_at=time.time(),
+                ttl=ttl if ttl is not None else self._default_ttl,
+            )
+
+    async def delete(self, key: str) -> bool:
+        """Delete a value from the cache."""
+        async with self._lock:
+            if key in self._cache:
+                del self._cache[key]
+                return True
+            return False
+
+    async def clear(self) -> None:
+        """Clear all cache entries."""
+        async with self._lock:
+            self._cache.clear()
+
+    async def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        async with self._lock:
+            entry = self._cache.get(key)
+            if entry is None:
+                return False
+            if entry.is_expired:
+                del self._cache[key]
+                return False
+            return True
+
+    def _evict_one(self) -> None:
+        """Evict one entry (LRU-like based on hits)."""
+        if not self._cache:
+            return
+
+        # First, remove expired entries
+        expired = [k for k, v in self._cache.items() if v.is_expired]
+        if expired:
+            del self._cache[expired[0]]
+            return
+
+        # Otherwise, remove entry with lowest hits
+        min_hits_key = min(self._cache.keys(), key=lambda k: self._cache[k].hits)
+        del self._cache[min_hits_key]
+
+    @property
+    def size(self) -> int:
+        """Get current cache size."""
+        return len(self._cache)
+
+
+class DiskCache(CacheBackend):
+    """Disk-based cache backend with TTL support.
+
+    Stores cached values as JSON files on disk.
+
+    Example:
+        >>> cache = DiskCache(path="/tmp/ai_cache", default_ttl=86400)
+        >>> await cache.set("key", {"data": "value"})
+        >>> value = await cache.get("key")
+    """
+
+    def __init__(
+        self,
+        path: str | Path,
+        default_ttl: float | None = None,
+        max_size_bytes: int = 100 * 1024 * 1024,  # 100MB default
+    ) -> None:
+        """Initialize disk cache.
+
+        Args:
+            path: Cache directory path
+            default_ttl: Default TTL in seconds
+            max_size_bytes: Maximum cache size in bytes
+        """
+        self._path = Path(path)
+        self._default_ttl = default_ttl
+        self._max_size_bytes = max_size_bytes
+        self._lock = asyncio.Lock()
+
+        # Create cache directory
+        self._path.mkdir(parents=True, exist_ok=True)
+
+    def _key_to_path(self, key: str) -> Path:
+        """Convert cache key to file path.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            File path
+        """
+        # Hash key to create safe filename
+        key_hash = hashlib.sha256(key.encode()).hexdigest()
+        return self._path / f"{key_hash}.json"
+
+    async def get(self, key: str) -> Any | None:
+        """Get a value from the cache."""
+        path = self._key_to_path(key)
+
+        async with self._lock:
+            if not path.exists():
+                return None
+
+            try:
+                content = path.read_text()
+                data = json.loads(content)
+
+                # Check expiry
+                if data.get("ttl") is not None:
+                    expires_at = data["created_at"] + data["ttl"]
+                    if time.time() > expires_at:
+                        path.unlink(missing_ok=True)
+                        return None
+
+                return data["value"]
+
+            except (json.JSONDecodeError, KeyError, OSError):
+                return None
+
+    async def set(
+        self, key: str, value: Any, ttl: float | None = None
+    ) -> None:
+        """Set a value in the cache."""
+        path = self._key_to_path(key)
+
+        async with self._lock:
+            data = {
+                "key": key,
+                "value": value,
+                "created_at": time.time(),
+                "ttl": ttl if ttl is not None else self._default_ttl,
+            }
+
+            try:
+                path.write_text(json.dumps(data, default=str))
+            except OSError:
+                # Ignore write errors
+                pass
+
+    async def delete(self, key: str) -> bool:
+        """Delete a value from the cache."""
+        path = self._key_to_path(key)
+
+        async with self._lock:
+            if path.exists():
+                path.unlink(missing_ok=True)
+                return True
+            return False
+
+    async def clear(self) -> None:
+        """Clear all cache entries."""
+        async with self._lock:
+            for path in self._path.glob("*.json"):
+                path.unlink(missing_ok=True)
+
+    async def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        path = self._key_to_path(key)
+        return path.exists()
+
+    async def cleanup_expired(self) -> int:
+        """Remove expired entries.
+
+        Returns:
+            Number of entries removed
+        """
+        removed = 0
+        now = time.time()
+
+        async with self._lock:
+            for path in self._path.glob("*.json"):
+                try:
+                    content = path.read_text()
+                    data = json.loads(content)
+
+                    if data.get("ttl") is not None:
+                        expires_at = data["created_at"] + data["ttl"]
+                        if now > expires_at:
+                            path.unlink(missing_ok=True)
+                            removed += 1
+
+                except (json.JSONDecodeError, KeyError, OSError):
+                    pass
+
+        return removed
+
+    @property
+    def cache_size_bytes(self) -> int:
+        """Get total cache size in bytes."""
+        total = 0
+        for path in self._path.glob("*.json"):
+            total += path.stat().st_size
+        return total
+
+
+class NullCache(CacheBackend):
+    """Null cache backend that doesn't cache anything.
+
+    Useful for testing or disabling caching.
+    """
+
+    async def get(self, key: str) -> Any | None:  # noqa: ARG002
+        """Always returns None."""
+        return None
+
+    async def set(
+        self, key: str, value: Any, ttl: float | None = None
+    ) -> None:
+        """Does nothing."""
+        pass
+
+    async def delete(self, key: str) -> bool:  # noqa: ARG002
+        """Always returns False."""
+        return False
+
+    async def clear(self) -> None:
+        """Does nothing."""
+        pass
+
+    async def exists(self, key: str) -> bool:  # noqa: ARG002
+        """Always returns False."""
+        return False

ai_lib_python/cache/key.py

@@ -0,0 +1,237 @@
+"""
+Cache key generation utilities.
+
+Provides deterministic cache key generation for requests.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class CacheKey:
+    """A cache key with metadata.
+
+    Attributes:
+        key: The cache key string
+        model: Model used
+        messages_hash: Hash of messages
+        params_hash: Hash of parameters
+    """
+
+    key: str
+    model: str = ""
+    messages_hash: str = ""
+    params_hash: str = ""
+
+    def __str__(self) -> str:
+        """Return the key string."""
+        return self.key
+
+    def __hash__(self) -> int:
+        """Return hash of the key."""
+        return hash(self.key)
+
+    def __eq__(self, other: object) -> bool:
+        """Check equality."""
+        if isinstance(other, CacheKey):
+            return self.key == other.key
+        if isinstance(other, str):
+            return self.key == other
+        return False
+
+
+class CacheKeyGenerator:
+    """Generates deterministic cache keys for requests.
+
+    Example:
+        >>> generator = CacheKeyGenerator()
+        >>> key = generator.generate(
+        ...     model="gpt-4o",
+        ...     messages=[{"role": "user", "content": "Hello"}],
+        ...     temperature=0.7,
+        ... )
+        >>> print(key.key)  # "ai:gpt-4o:a1b2c3..."
+    """
+
+    def __init__(
+        self,
+        prefix: str = "ai",
+        include_model: bool = True,
+        include_params: bool = True,
+        excluded_params: list[str] | None = None,
+    ) -> None:
+        """Initialize key generator.
+
+        Args:
+            prefix: Key prefix
+            include_model: Whether to include model in key
+            include_params: Whether to include params in key
+            excluded_params: Parameters to exclude from key
+        """
+        self._prefix = prefix
+        self._include_model = include_model
+        self._include_params = include_params
+        self._excluded_params = set(excluded_params or ["user", "stream"])
+
+    def generate(
+        self,
+        model: str,
+        messages: list[dict[str, Any]],
+        **params: Any,
+    ) -> CacheKey:
+        """Generate a cache key.
+
+        Args:
+            model: Model name
+            messages: Chat messages
+            **params: Additional parameters
+
+        Returns:
+            CacheKey instance
+        """
+        # Hash messages
+        messages_hash = self._hash_messages(messages)
+
+        # Hash parameters
+        filtered_params = {
+            k: v for k, v in params.items() if k not in self._excluded_params
+        }
+        params_hash = self._hash_params(filtered_params)
+
+        # Build key
+        parts = [self._prefix]
+
+        if self._include_model:
+            parts.append(model)
+
+        parts.append(messages_hash[:16])
+
+        if self._include_params and params_hash:
+            parts.append(params_hash[:8])
+
+        key = ":".join(parts)
+
+        return CacheKey(
+            key=key,
+            model=model,
+            messages_hash=messages_hash,
+            params_hash=params_hash,
+        )
+
+    def generate_for_embedding(
+        self,
+        model: str,
+        input_text: str | list[str],
+        dimensions: int | None = None,
+    ) -> CacheKey:
+        """Generate cache key for embedding request.
+
+        Args:
+            model: Model name
+            input_text: Input text or list of texts
+            dimensions: Output dimensions
+
+        Returns:
+            CacheKey instance
+        """
+        # Normalize input
+        if isinstance(input_text, str):
+            input_hash = self._hash_string(input_text)
+        else:
+            input_hash = self._hash_string(json.dumps(input_text, sort_keys=True))
+
+        # Build key
+        parts = [self._prefix, "emb", model, input_hash[:16]]
+
+        if dimensions:
+            parts.append(str(dimensions))
+
+        key = ":".join(parts)
+
+        return CacheKey(
+            key=key,
+            model=model,
+            messages_hash=input_hash,
+        )
+
+    def _hash_messages(self, messages: list[dict[str, Any]]) -> str:
+        """Hash a list of messages.
+
+        Args:
+            messages: Messages to hash
+
+        Returns:
+            Hash string
+        """
+        # Normalize messages for hashing
+        normalized = []
+        for msg in messages:
+            normalized.append({
+                "role": msg.get("role", ""),
+                "content": self._normalize_content(msg.get("content", "")),
+            })
+
+        content = json.dumps(normalized, sort_keys=True, ensure_ascii=True)
+        return self._hash_string(content)
+
+    def _normalize_content(self, content: Any) -> Any:
+        """Normalize message content for hashing.
+
+        Args:
+            content: Content to normalize
+
+        Returns:
+            Normalized content
+        """
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            # Handle content blocks
+            normalized = []
+            for block in content:
+                if isinstance(block, dict):
+                    if block.get("type") == "text":
+                        normalized.append({"type": "text", "text": block.get("text", "")})
+                    elif block.get("type") == "image_url":
+                        # Include image URL in hash
+                        normalized.append({
+                            "type": "image_url",
+                            "url": block.get("image_url", {}).get("url", ""),
+                        })
+                    else:
+                        normalized.append(block)
+                else:
+                    normalized.append(block)
+            return normalized
+        return content
+
+    def _hash_params(self, params: dict[str, Any]) -> str:
+        """Hash parameters.
+
+        Args:
+            params: Parameters to hash
+
+        Returns:
+            Hash string
+        """
+        if not params:
+            return ""
+
+        content = json.dumps(params, sort_keys=True, ensure_ascii=True)
+        return self._hash_string(content)
+
+    def _hash_string(self, content: str) -> str:
+        """Hash a string using SHA-256.
+
+        Args:
+            content: String to hash
+
+        Returns:
+            Hex digest
+        """
+        return hashlib.sha256(content.encode()).hexdigest()