flashlite 0.1.0__py3-none-any.whl

flashlite/__init__.py

@@ -0,0 +1,169 @@
+"""
+Flashlite - Batteries-included wrapper for litellm.
+
+Features:
+- Rate limiting with token bucket algorithm
+- Retries with exponential backoff
+- Jinja templating for prompts
+- Async-first with sync wrappers
+- Full passthrough of provider kwargs
+- Response caching (memory and disk backends)
+- Cost tracking and budget limits
+- Structured logging and Inspect framework integration
+"""
+
+from .cache import CacheBackend, DiskCache, MemoryCache, generate_cache_key
+from .client import Flashlite
+from .config import FlashliteConfig, load_env_files, validate_api_keys
+from .conversation import (
+    Agent,
+    ChatMessage,
+    ContextLimits,
+    ContextManager,
+    Conversation,
+    ConversationState,
+    MultiAgentChat,
+    Turn,
+    estimate_messages_tokens,
+    estimate_tokens,
+    truncate_messages,
+)
+from .core.messages import (
+    assistant_message,
+    format_messages,
+    system_message,
+    tool_message,
+    user_message,
+)
+from .observability import (
+    BudgetExceededError,
+    CallbackManager,
+    CostTracker,
+    FlashliteModelAPI,
+    InspectLogger,
+    StructuredLogger,
+)
+from .structured import (
+    StructuredOutputError,
+    generate_json_schema,
+    parse_json_response,
+    schema_to_prompt,
+)
+from .templating import TemplateEngine, TemplateRegistry
+from .tools import (
+    ToolCall,
+    ToolDefinition,
+    ToolLoopResult,
+    ToolRegistry,
+    ToolResult,
+    format_tool_result,
+    run_tool_loop,
+    tool,
+    tool_from_pydantic,
+    tools_to_anthropic,
+    tools_to_openai,
+)
+from .types import (
+    CompletionError,
+    # Request/Response types
+    CompletionRequest,
+    CompletionResponse,
+    ConfigError,
+    # Exceptions
+    FlashliteError,
+    # Message types
+    Message,
+    MessageDict,
+    Messages,
+    RateLimitConfig,
+    RateLimitError,
+    # Configuration types
+    RetryConfig,
+    Role,
+    TemplateError,
+    ThinkingConfig,
+    UsageInfo,
+    ValidationError,
+    thinking_enabled,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Main client
+    "Flashlite",
+    # Configuration
+    "FlashliteConfig",
+    "RetryConfig",
+    "RateLimitConfig",
+    "ThinkingConfig",
+    "thinking_enabled",
+    "load_env_files",
+    "validate_api_keys",
+    # Request/Response
+    "CompletionRequest",
+    "CompletionResponse",
+    "UsageInfo",
+    # Messages
+    "Message",
+    "Messages",
+    "MessageDict",
+    "Role",
+    "format_messages",
+    "user_message",
+    "system_message",
+    "assistant_message",
+    "tool_message",
+    # Templating
+    "TemplateEngine",
+    "TemplateRegistry",
+    # Caching
+    "CacheBackend",
+    "MemoryCache",
+    "DiskCache",
+    "generate_cache_key",
+    # Conversation
+    "Conversation",
+    "ConversationState",
+    "Turn",
+    "ContextManager",
+    "ContextLimits",
+    "estimate_tokens",
+    "estimate_messages_tokens",
+    "truncate_messages",
+    # Multi-agent
+    "MultiAgentChat",
+    "Agent",
+    "ChatMessage",
+    # Observability
+    "StructuredLogger",
+    "CostTracker",
+    "BudgetExceededError",
+    "CallbackManager",
+    "InspectLogger",
+    "FlashliteModelAPI",
+    # Structured outputs
+    "StructuredOutputError",
+    "generate_json_schema",
+    "schema_to_prompt",
+    "parse_json_response",
+    # Exceptions
+    "FlashliteError",
+    "CompletionError",
+    "RateLimitError",
+    "ValidationError",
+    "TemplateError",
+    "ConfigError",
+    # Tools
+    "tool",
+    "tool_from_pydantic",
+    "ToolDefinition",
+    "ToolRegistry",
+    "ToolCall",
+    "ToolResult",
+    "ToolLoopResult",
+    "run_tool_loop",
+    "tools_to_openai",
+    "tools_to_anthropic",
+    "format_tool_result",
+]

flashlite/cache/__init__.py

@@ -0,0 +1,14 @@
+"""Caching module for flashlite."""
+
+from .base import CacheBackend, CacheEntry, generate_cache_key, is_cacheable_request
+from .disk import DiskCache
+from .memory import MemoryCache
+
+__all__ = [
+    "CacheBackend",
+    "CacheEntry",
+    "generate_cache_key",
+    "is_cacheable_request",
+    "MemoryCache",
+    "DiskCache",
+]

flashlite/cache/base.py

@@ -0,0 +1,194 @@
+"""Base cache protocol and key generation for flashlite."""
+
+import hashlib
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+from ..types import CompletionRequest, CompletionResponse
+
+
+@dataclass
+class CacheEntry:
+    """A cached completion response with metadata."""
+
+    response: CompletionResponse
+    request_hash: str
+    created_at: float
+    ttl: float | None = None  # None means no expiration
+
+    def is_expired(self, current_time: float) -> bool:
+        """Check if this cache entry has expired."""
+        if self.ttl is None:
+            return False
+        return current_time > self.created_at + self.ttl
+
+
+class CacheBackend(ABC):
+    """Abstract base class for cache backends."""
+
+    @abstractmethod
+    async def get(self, key: str) -> CompletionResponse | None:
+        """
+        Retrieve a cached response.
+
+        Args:
+            key: The cache key
+
+        Returns:
+            The cached CompletionResponse, or None if not found/expired
+        """
+        pass
+
+    @abstractmethod
+    async def set(
+        self,
+        key: str,
+        response: CompletionResponse,
+        ttl: float | None = None,
+    ) -> None:
+        """
+        Store a response in the cache.
+
+        Args:
+            key: The cache key
+            response: The response to cache
+            ttl: Time-to-live in seconds (None = no expiration)
+        """
+        pass
+
+    @abstractmethod
+    async def delete(self, key: str) -> bool:
+        """
+        Delete a cached entry.
+
+        Args:
+            key: The cache key
+
+        Returns:
+            True if the key existed and was deleted
+        """
+        pass
+
+    @abstractmethod
+    async def clear(self) -> int:
+        """
+        Clear all cached entries.
+
+        Returns:
+            Number of entries cleared
+        """
+        pass
+
+    @abstractmethod
+    async def size(self) -> int:
+        """
+        Get the number of cached entries.
+
+        Returns:
+            Number of entries in the cache
+        """
+        pass
+
+    async def close(self) -> None:
+        """Close any resources held by the cache backend."""
+        pass
+
+
+def generate_cache_key(request: CompletionRequest) -> str:
+    """
+    Generate a deterministic cache key for a completion request.
+
+    The key is based on:
+    - Model name
+    - Messages (serialized)
+    - Temperature (if set)
+    - Other deterministic parameters
+
+    Note: Requests with temperature > 0 or reasoning enabled are typically
+    not good candidates for caching since responses are non-deterministic.
+
+    Args:
+        request: The completion request
+
+    Returns:
+        A hex-encoded SHA-256 hash as the cache key
+    """
+    # Build the key components
+    key_data: dict[str, Any] = {
+        "model": request.model,
+        "messages": [dict(m) for m in request.messages],
+    }
+
+    # Include parameters that affect output
+    if request.temperature is not None:
+        key_data["temperature"] = request.temperature
+    if request.max_tokens is not None:
+        key_data["max_tokens"] = request.max_tokens
+    if request.max_completion_tokens is not None:
+        key_data["max_completion_tokens"] = request.max_completion_tokens
+    if request.top_p is not None:
+        key_data["top_p"] = request.top_p
+    if request.stop is not None:
+        key_data["stop"] = request.stop
+    if request.reasoning_effort is not None:
+        key_data["reasoning_effort"] = request.reasoning_effort
+    if request.thinking is not None:
+        key_data["thinking"] = request.thinking
+
+    # Include extra kwargs that affect output (excluding metadata-only ones and test helpers)
+    exclude_from_key = {"timeout", "metadata", "tags", "user", "mock_response"}
+    for k, v in request.extra_kwargs.items():
+        if k not in exclude_from_key:
+            key_data[f"extra.{k}"] = v
+
+    # Serialize to JSON with sorted keys for determinism
+    serialized = json.dumps(key_data, sort_keys=True, default=str)
+
+    # Hash to get fixed-length key
+    return hashlib.sha256(serialized.encode()).hexdigest()
+
+
+def is_cacheable_request(request: CompletionRequest) -> tuple[bool, str | None]:
+    """
+    Check if a request is suitable for caching.
+
+    Returns:
+        Tuple of (is_cacheable, warning_message).
+        If is_cacheable is True, warning_message may still contain a warning.
+        If is_cacheable is False, warning_message explains why.
+    """
+    warnings: list[str] = []
+
+    # Check for non-deterministic temperature
+    if request.temperature is not None and request.temperature > 0:
+        warnings.append(
+            f"temperature={request.temperature} > 0 may produce different outputs"
+        )
+
+    # Check for reasoning models (non-deterministic by nature)
+    if request.reasoning_effort is not None:
+        warnings.append(
+            f"reasoning_effort='{request.reasoning_effort}' - "
+            "reasoning models may produce varying outputs"
+        )
+
+    if request.thinking is not None:
+        warnings.append(
+            "thinking enabled - extended thinking models may produce varying outputs"
+        )
+
+    # Check model name for known reasoning models
+    model_lower = request.model.lower()
+    reasoning_model_patterns = ["o1", "o3", "claude-3-5-sonnet", "claude-sonnet-4"]
+    if any(pattern in model_lower for pattern in reasoning_model_patterns):
+        if not any("reasoning" in w for w in warnings):
+            warnings.append(
+                f"model '{request.model}' appears to be a reasoning model - outputs may vary"
+            )
+
+    if warnings:
+        return True, "; ".join(warnings)
+
+    return True, None

flashlite/cache/disk.py

@@ -0,0 +1,285 @@
+"""SQLite-based disk cache for persistent caching."""
+
+import asyncio
+import json
+import sqlite3
+import time
+from pathlib import Path
+from typing import Any
+
+from ..types import CompletionResponse, UsageInfo
+from .base import CacheBackend
+
+
+class DiskCache(CacheBackend):
+    """
+    SQLite-based disk cache for persistent caching.
+
+    This cache stores responses in a SQLite database, providing
+    persistence across process restarts.
+
+    Example:
+        cache = DiskCache("./cache/completions.db", default_ttl=86400)
+
+        # Store a response
+        await cache.set(key, response)
+
+        # Retrieve (returns None if expired or not found)
+        cached = await cache.get(key)
+
+        # Close when done
+        await cache.close()
+    """
+
+    def __init__(
+        self,
+        path: str | Path,
+        default_ttl: float | None = None,
+        auto_vacuum: bool = True,
+    ):
+        """
+        Initialize the disk cache.
+
+        Args:
+            path: Path to SQLite database file (will be created if doesn't exist)
+            default_ttl: Default time-to-live in seconds (None = no expiration)
+            auto_vacuum: Whether to run VACUUM on startup to reclaim space
+        """
+        self._path = Path(path)
+        self._default_ttl = default_ttl
+        self._auto_vacuum = auto_vacuum
+        self._conn: sqlite3.Connection | None = None
+        self._lock = asyncio.Lock()
+        self._hits = 0
+        self._misses = 0
+
+        # Ensure parent directory exists
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Initialize database
+        self._init_db()
+
+    def _init_db(self) -> None:
+        """Initialize the database schema."""
+        self._conn = sqlite3.connect(str(self._path), check_same_thread=False)
+        self._conn.row_factory = sqlite3.Row
+
+        # Create tables
+        self._conn.execute("""
+            CREATE TABLE IF NOT EXISTS cache (
+                key TEXT PRIMARY KEY,
+                response_json TEXT NOT NULL,
+                created_at REAL NOT NULL,
+                expires_at REAL,
+                model TEXT,
+                input_tokens INTEGER,
+                output_tokens INTEGER
+            )
+        """)
+
+        # Create index for expiration queries
+        self._conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_expires_at ON cache(expires_at)
+        """)
+
+        self._conn.commit()
+
+        # Optional vacuum to reclaim space
+        if self._auto_vacuum:
+            try:
+                self._conn.execute("VACUUM")
+            except sqlite3.OperationalError:
+                # VACUUM can fail if database is locked
+                pass
+
+    def _serialize_response(self, response: CompletionResponse) -> str:
+        """Serialize a CompletionResponse to JSON."""
+        data: dict[str, Any] = {
+            "content": response.content,
+            "model": response.model,
+            "finish_reason": response.finish_reason,
+        }
+
+        if response.usage:
+            data["usage"] = {
+                "input_tokens": response.usage.input_tokens,
+                "output_tokens": response.usage.output_tokens,
+                "total_tokens": response.usage.total_tokens,
+            }
+
+        return json.dumps(data)
+
+    def _deserialize_response(self, json_str: str) -> CompletionResponse:
+        """Deserialize a CompletionResponse from JSON."""
+        data = json.loads(json_str)
+
+        usage = None
+        if "usage" in data:
+            usage = UsageInfo(
+                input_tokens=data["usage"]["input_tokens"],
+                output_tokens=data["usage"]["output_tokens"],
+                total_tokens=data["usage"]["total_tokens"],
+            )
+
+        return CompletionResponse(
+            content=data["content"],
+            model=data["model"],
+            finish_reason=data.get("finish_reason"),
+            usage=usage,
+        )
+
+    async def get(self, key: str) -> CompletionResponse | None:
+        """Retrieve a cached response."""
+        async with self._lock:
+            if self._conn is None:
+                raise RuntimeError("Cache has been closed")
+
+            now = time.time()
+
+            # Query for non-expired entry
+            cursor = self._conn.execute(
+                """
+                SELECT response_json FROM cache
+                WHERE key = ? AND (expires_at IS NULL OR expires_at > ?)
+                """,
+                (key, now),
+            )
+            row = cursor.fetchone()
+
+            if row is None:
+                self._misses += 1
+                return None
+
+            self._hits += 1
+            return self._deserialize_response(row["response_json"])
+
+    async def set(
+        self,
+        key: str,
+        response: CompletionResponse,
+        ttl: float | None = None,
+    ) -> None:
+        """Store a response in the cache."""
+        async with self._lock:
+            if self._conn is None:
+                raise RuntimeError("Cache has been closed")
+
+            effective_ttl = ttl if ttl is not None else self._default_ttl
+            now = time.time()
+            expires_at = now + effective_ttl if effective_ttl else None
+
+            response_json = self._serialize_response(response)
+
+            self._conn.execute(
+                """
+                INSERT OR REPLACE INTO cache
+                (key, response_json, created_at, expires_at, model, input_tokens, output_tokens)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    key,
+                    response_json,
+                    now,
+                    expires_at,
+                    response.model,
+                    response.usage.input_tokens if response.usage else None,
+                    response.usage.output_tokens if response.usage else None,
+                ),
+            )
+            self._conn.commit()
+
+    async def delete(self, key: str) -> bool:
+        """Delete a cached entry."""
+        async with self._lock:
+            if self._conn is None:
+                raise RuntimeError("Cache has been closed")
+
+            cursor = self._conn.execute("DELETE FROM cache WHERE key = ?", (key,))
+            self._conn.commit()
+            return cursor.rowcount > 0
+
+    async def clear(self) -> int:
+        """Clear all cached entries."""
+        async with self._lock:
+            if self._conn is None:
+                raise RuntimeError("Cache has been closed")
+
+            cursor = self._conn.execute("SELECT COUNT(*) as count FROM cache")
+            count = cursor.fetchone()["count"]
+
+            self._conn.execute("DELETE FROM cache")
+            self._conn.commit()
+
+            self._hits = 0
+            self._misses = 0
+
+            return count
+
+    async def size(self) -> int:
+        """Get the number of cached entries (excluding expired)."""
+        async with self._lock:
+            if self._conn is None:
+                raise RuntimeError("Cache has been closed")
+
+            now = time.time()
+            cursor = self._conn.execute(
+                """
+                SELECT COUNT(*) as count FROM cache
+                WHERE expires_at IS NULL OR expires_at > ?
+                """,
+                (now,),
+            )
+            return cursor.fetchone()["count"]
+
+    async def cleanup_expired(self) -> int:
+        """
+        Remove expired entries from the cache.
+
+        Returns:
+            Number of entries removed
+        """
+        async with self._lock:
+            if self._conn is None:
+                raise RuntimeError("Cache has been closed")
+
+            now = time.time()
+            cursor = self._conn.execute(
+                "DELETE FROM cache WHERE expires_at IS NOT NULL AND expires_at <= ?",
+                (now,),
+            )
+            self._conn.commit()
+            return cursor.rowcount
+
+    async def close(self) -> None:
+        """Close the database connection."""
+        async with self._lock:
+            if self._conn is not None:
+                self._conn.close()
+                self._conn = None
+
+    @property
+    def hits(self) -> int:
+        """Number of cache hits (this session only)."""
+        return self._hits
+
+    @property
+    def misses(self) -> int:
+        """Number of cache misses (this session only)."""
+        return self._misses
+
+    @property
+    def hit_rate(self) -> float:
+        """Cache hit rate (0.0 to 1.0) for this session."""
+        total = self._hits + self._misses
+        if total == 0:
+            return 0.0
+        return self._hits / total
+
+    def stats(self) -> dict[str, Any]:
+        """Get cache statistics."""
+        return {
+            "path": str(self._path),
+            "hits": self._hits,
+            "misses": self._misses,
+            "hit_rate": self.hit_rate,
+        }