flashlite 0.1.0__py3-none-any.whl

flashlite/middleware/base.py

@@ -0,0 +1,90 @@
+"""Base middleware protocol and chain implementation."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Awaitable, Callable
+
+from ..types import CompletionRequest, CompletionResponse
+
+# Type alias for the completion handler
+CompletionHandler = Callable[[CompletionRequest], Awaitable[CompletionResponse]]
+
+
+class Middleware(ABC):
+    """
+    Abstract base class for middleware.
+
+    Middleware can intercept requests before they're sent and responses
+    after they're received. They form a chain where each middleware
+    calls the next one (or short-circuits).
+    """
+
+    @abstractmethod
+    async def __call__(
+        self,
+        request: CompletionRequest,
+        next_handler: CompletionHandler,
+    ) -> CompletionResponse:
+        """
+        Process a request.
+
+        Args:
+            request: The completion request
+            next_handler: The next middleware or final handler to call
+
+        Returns:
+            The completion response
+        """
+        pass
+
+
+class MiddlewareChain:
+    """
+    Chains multiple middleware together.
+
+    Middleware are executed in order, with each one wrapping the next.
+    The final handler (actual API call) is at the end of the chain.
+    """
+
+    def __init__(
+        self,
+        middleware: list[Middleware],
+        final_handler: CompletionHandler,
+    ):
+        """
+        Initialize the middleware chain.
+
+        Args:
+            middleware: List of middleware to apply (in order)
+            final_handler: The final handler (actual completion call)
+        """
+        self._middleware = middleware
+        self._final_handler = final_handler
+
+    async def __call__(self, request: CompletionRequest) -> CompletionResponse:
+        """Execute the middleware chain."""
+        return await self._execute(request, 0)
+
+    async def _execute(self, request: CompletionRequest, index: int) -> CompletionResponse:
+        """Recursively execute middleware chain."""
+        if index >= len(self._middleware):
+            # No more middleware, call the final handler
+            return await self._final_handler(request)
+
+        # Get current middleware and create next handler
+        current = self._middleware[index]
+
+        async def next_handler(req: CompletionRequest) -> CompletionResponse:
+            return await self._execute(req, index + 1)
+
+        return await current(request, next_handler)
+
+
+class PassthroughMiddleware(Middleware):
+    """A middleware that does nothing - useful for testing."""
+
+    async def __call__(
+        self,
+        request: CompletionRequest,
+        next_handler: CompletionHandler,
+    ) -> CompletionResponse:
+        return await next_handler(request)

flashlite/middleware/cache.py

@@ -0,0 +1,121 @@
+"""Cache middleware for flashlite."""
+
+import logging
+from collections.abc import Awaitable, Callable
+
+from ..cache.base import CacheBackend, generate_cache_key, is_cacheable_request
+from ..types import CompletionRequest, CompletionResponse
+from .base import Middleware
+
+logger = logging.getLogger(__name__)
+
+
+class CacheMiddleware(Middleware):
+    """
+    Middleware that caches completion responses.
+
+    Caches responses based on a hash of the request parameters.
+    Emits warnings when caching is used with non-deterministic settings
+    (temperature > 0 or reasoning models).
+
+    Example:
+        cache = MemoryCache(max_size=1000)
+        middleware = CacheMiddleware(cache)
+    """
+
+    def __init__(
+        self,
+        backend: CacheBackend,
+        ttl: float | None = None,
+        warn_non_deterministic: bool = True,
+    ):
+        """
+        Initialize the cache middleware.
+
+        Args:
+            backend: The cache backend to use
+            ttl: Default TTL for cached entries (seconds)
+            warn_non_deterministic: Whether to warn about non-deterministic caching
+        """
+        self._backend = backend
+        self._ttl = ttl
+        self._warn_non_deterministic = warn_non_deterministic
+        self._warned_keys: set[str] = set()  # Track warnings to avoid spam
+
+    async def __call__(
+        self,
+        request: CompletionRequest,
+        next_handler: Callable[[CompletionRequest], Awaitable[CompletionResponse]],
+    ) -> CompletionResponse:
+        """Process request with caching."""
+        # Generate cache key
+        cache_key = generate_cache_key(request)
+
+        # Check if request is suitable for caching and emit warnings
+        if self._warn_non_deterministic:
+            _, warning = is_cacheable_request(request)
+            if warning and cache_key not in self._warned_keys:
+                logger.warning(
+                    f"Caching enabled but request may be non-deterministic: {warning}. "
+                    "Consider disabling cache for this request with force_refresh=True, "
+                    "or set temperature=0 for deterministic outputs."
+                )
+                self._warned_keys.add(cache_key)
+
+        # Try to get from cache
+        cached_response = await self._backend.get(cache_key)
+        if cached_response is not None:
+            logger.debug(f"Cache hit for key {cache_key[:16]}...")
+            return cached_response
+
+        # Cache miss - call the next handler
+        logger.debug(f"Cache miss for key {cache_key[:16]}...")
+        response = await next_handler(request)
+
+        # Store in cache
+        await self._backend.set(cache_key, response, self._ttl)
+
+        return response
+
+    @property
+    def backend(self) -> CacheBackend:
+        """Get the cache backend."""
+        return self._backend
+
+
+class CacheConfig:
+    """
+    Configuration for caching behavior.
+
+    Note: Caching is disabled by default. When enabled, warnings are emitted
+    for non-deterministic requests (temperature > 0 or reasoning models).
+    """
+
+    def __init__(
+        self,
+        enabled: bool = False,
+        backend: CacheBackend | None = None,
+        ttl: float | None = None,
+        warn_non_deterministic: bool = True,
+    ):
+        """
+        Initialize cache configuration.
+
+        Args:
+            enabled: Whether caching is enabled (default: False)
+            backend: The cache backend to use
+            ttl: Default TTL for cached entries (seconds)
+            warn_non_deterministic: Whether to warn about non-deterministic caching
+        """
+        self.enabled = enabled
+        self.backend = backend
+        self.ttl = ttl
+        self.warn_non_deterministic = warn_non_deterministic
+
+        # Emit info message about caching status
+        if not enabled:
+            logger.info(
+                "Caching is disabled by default. To enable, pass "
+                "cache=CacheConfig(enabled=True, backend=...) or "
+                "cache=MemoryCache(...) to the Flashlite client."
+            )

flashlite/middleware/logging.py

@@ -0,0 +1,159 @@
+"""Logging middleware for flashlite."""
+
+import logging
+import time
+import uuid
+from collections.abc import Awaitable, Callable
+
+from ..observability.callbacks import CallbackManager
+from ..observability.logging import StructuredLogger
+from ..observability.metrics import CostTracker
+from ..types import CompletionRequest, CompletionResponse
+from .base import Middleware
+
+logger = logging.getLogger(__name__)
+
+
+class LoggingMiddleware(Middleware):
+    """
+    Middleware that logs requests and responses.
+
+    Supports structured logging to files, callback-based logging,
+    and cost tracking.
+
+    Example:
+        structured_logger = StructuredLogger(log_file="./logs/completions.jsonl")
+        cost_tracker = CostTracker(budget_limit=10.0)
+
+        middleware = LoggingMiddleware(
+            logger=structured_logger,
+            cost_tracker=cost_tracker,
+            log_level="INFO",
+        )
+    """
+
+    def __init__(
+        self,
+        structured_logger: StructuredLogger | None = None,
+        cost_tracker: CostTracker | None = None,
+        callbacks: CallbackManager | None = None,
+        log_level: str = "INFO",
+    ):
+        """
+        Initialize the logging middleware.
+
+        Args:
+            structured_logger: Structured logger for file logging
+            cost_tracker: Cost tracker for budget monitoring
+            callbacks: Callback manager for event hooks
+            log_level: Minimum log level for standard logging
+        """
+        self._structured_logger = structured_logger
+        self._cost_tracker = cost_tracker
+        self._callbacks = callbacks
+        self._log_level = getattr(logging, log_level.upper())
+
+    async def __call__(
+        self,
+        request: CompletionRequest,
+        next_handler: Callable[[CompletionRequest], Awaitable[CompletionResponse]],
+    ) -> CompletionResponse:
+        """Process request with logging."""
+        request_id = str(uuid.uuid4())
+        start_time = time.perf_counter()
+
+        # Log request
+        if self._structured_logger:
+            self._structured_logger.log_request(request, request_id)
+
+        if self._callbacks:
+            await self._callbacks.emit_request(request, request_id)
+
+        logger.log(
+            self._log_level,
+            f"[{request_id[:8]}] Starting request: model={request.model}",
+        )
+
+        try:
+            # Call next handler
+            response = await next_handler(request)
+
+            # Calculate latency
+            latency_ms = (time.perf_counter() - start_time) * 1000
+
+            # Log response
+            if self._structured_logger:
+                self._structured_logger.log_response(
+                    response, request_id, latency_ms, cached=False
+                )
+
+            if self._callbacks:
+                await self._callbacks.emit_response(
+                    response, request_id, latency_ms, cached=False
+                )
+
+            # Track cost
+            if self._cost_tracker:
+                cost = self._cost_tracker.track(response)
+                logger.debug(
+                    f"[{request_id[:8]}] Cost: ${cost:.6f}, "
+                    f"Total: ${self._cost_tracker.total_cost:.4f}"
+                )
+
+            logger.log(
+                self._log_level,
+                f"[{request_id[:8]}] Completed: {latency_ms:.1f}ms, "
+                f"tokens={response.usage.total_tokens if response.usage else 'N/A'}",
+            )
+
+            return response
+
+        except Exception as e:
+            latency_ms = (time.perf_counter() - start_time) * 1000
+
+            if self._structured_logger:
+                self._structured_logger.log_error(request_id, e, latency_ms)
+
+            if self._callbacks:
+                await self._callbacks.emit_error(e, request_id, latency_ms)
+
+            logger.error(
+                f"[{request_id[:8]}] Error after {latency_ms:.1f}ms: {e}",
+            )
+            raise
+
+
+class CostTrackingMiddleware(Middleware):
+    """
+    Lightweight middleware that only tracks costs.
+
+    Use this when you want cost tracking without full logging.
+
+    Example:
+        tracker = CostTracker(budget_limit=10.0)
+        middleware = CostTrackingMiddleware(tracker)
+    """
+
+    def __init__(self, cost_tracker: CostTracker):
+        """
+        Initialize the cost tracking middleware.
+
+        Args:
+            cost_tracker: The cost tracker to use
+        """
+        self._cost_tracker = cost_tracker
+
+    async def __call__(
+        self,
+        request: CompletionRequest,
+        next_handler: Callable[[CompletionRequest], Awaitable[CompletionResponse]],
+    ) -> CompletionResponse:
+        """Process request with cost tracking."""
+        response = await next_handler(request)
+        self._cost_tracker.track(response)
+        return response
+
+    @property
+    def tracker(self) -> CostTracker:
+        """Get the cost tracker."""
+        return self._cost_tracker

flashlite/middleware/rate_limit.py

@@ -0,0 +1,211 @@
+"""Rate limiting middleware using token bucket algorithm."""
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+
+from ..types import CompletionRequest, CompletionResponse, RateLimitConfig, RateLimitError
+from .base import CompletionHandler, Middleware
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TokenBucket:
+    """
+    Token bucket rate limiter.
+
+    Tokens are added at a constant rate up to a maximum capacity.
+    Each request consumes one or more tokens. If not enough tokens
+    are available, the request waits.
+    """
+
+    rate: float  # Tokens added per second
+    capacity: float  # Maximum tokens in bucket
+    tokens: float = field(init=False)
+    last_update: float = field(init=False)
+    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
+
+    def __post_init__(self) -> None:
+        self.tokens = self.capacity
+        self.last_update = time.monotonic()
+
+    def _refill(self) -> None:
+        """Refill tokens based on elapsed time."""
+        now = time.monotonic()
+        elapsed = now - self.last_update
+        self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
+        self.last_update = now
+
+    async def acquire(self, tokens: float = 1.0, timeout: float | None = None) -> float:
+        """
+        Acquire tokens from the bucket.
+
+        Args:
+            tokens: Number of tokens to acquire
+            timeout: Maximum time to wait (None = wait forever)
+
+        Returns:
+            Time waited in seconds
+
+        Raises:
+            RateLimitError: If timeout exceeded
+        """
+        start_time = time.monotonic()
+        deadline = start_time + timeout if timeout else None
+
+        async with self._lock:
+            while True:
+                self._refill()
+
+                if self.tokens >= tokens:
+                    self.tokens -= tokens
+                    return time.monotonic() - start_time
+
+                # Calculate wait time for enough tokens
+                tokens_needed = tokens - self.tokens
+                wait_time = tokens_needed / self.rate
+
+                # Check timeout
+                if deadline:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        raise RateLimitError(
+                            f"Rate limit timeout after {timeout}s",
+                            retry_after=wait_time,
+                        )
+                    wait_time = min(wait_time, remaining)
+
+                # Release lock while waiting
+                self._lock.release()
+                try:
+                    await asyncio.sleep(wait_time)
+                finally:
+                    await self._lock.acquire()
+
+    @property
+    def available_tokens(self) -> float:
+        """Get current available tokens (without acquiring lock)."""
+        elapsed = time.monotonic() - self.last_update
+        return min(self.capacity, self.tokens + elapsed * self.rate)
+
+
+class RateLimitMiddleware(Middleware):
+    """
+    Middleware that enforces rate limits using token bucket algorithm.
+
+    Supports both requests-per-minute (RPM) and tokens-per-minute (TPM) limits.
+    """
+
+    def __init__(self, config: RateLimitConfig | None = None):
+        """
+        Initialize rate limit middleware.
+
+        Args:
+            config: Rate limit configuration
+        """
+        self.config = config or RateLimitConfig()
+        self._rpm_bucket: TokenBucket | None = None
+        self._tpm_bucket: TokenBucket | None = None
+        self._initialized = False
+
+    def _ensure_initialized(self) -> None:
+        """Lazily initialize buckets."""
+        if self._initialized:
+            return
+
+        if self.config.requests_per_minute:
+            # Convert RPM to requests per second
+            rate = self.config.requests_per_minute / 60.0
+            # Capacity allows small bursts (10% of per-minute rate)
+            capacity = max(1.0, self.config.requests_per_minute * 0.1)
+            self._rpm_bucket = TokenBucket(rate=rate, capacity=capacity)
+            logger.debug(
+                f"Rate limiter initialized: {self.config.requests_per_minute} RPM "
+                f"(rate={rate:.2f}/s, capacity={capacity:.1f})"
+            )
+
+        if self.config.tokens_per_minute:
+            rate = self.config.tokens_per_minute / 60.0
+            capacity = max(1000.0, self.config.tokens_per_minute * 0.1)
+            self._tpm_bucket = TokenBucket(rate=rate, capacity=capacity)
+            logger.debug(
+                f"Token rate limiter initialized: {self.config.tokens_per_minute} TPM"
+            )
+
+        self._initialized = True
+
+    async def __call__(
+        self,
+        request: CompletionRequest,
+        next_handler: CompletionHandler,
+    ) -> CompletionResponse:
+        """Execute with rate limiting."""
+        self._ensure_initialized()
+
+        # Acquire RPM token before making request
+        if self._rpm_bucket:
+            wait_time = await self._rpm_bucket.acquire()
+            if wait_time > 0.1:  # Only log significant waits
+                logger.debug(f"Rate limit: waited {wait_time:.2f}s for RPM token")
+
+        # Make the request
+        response = await next_handler(request)
+
+        # For TPM limiting, consume tokens based on actual usage
+        # This is post-hoc - we can't know token count before the request
+        if self._tpm_bucket and response.usage:
+            total_tokens = response.usage.total_tokens
+            if total_tokens > 0:
+                # Don't block on TPM - just record the usage
+                # This creates backpressure for subsequent requests
+                await self._tpm_bucket.acquire(tokens=float(total_tokens))
+
+        return response
+
+    @property
+    def rpm_available(self) -> float | None:
+        """Get available RPM tokens."""
+        if self._rpm_bucket:
+            return self._rpm_bucket.available_tokens
+        return None
+
+    @property
+    def tpm_available(self) -> float | None:
+        """Get available TPM tokens."""
+        if self._tpm_bucket:
+            return self._tpm_bucket.available_tokens
+        return None
+
+
+class ConcurrencyLimiter:
+    """
+    Limits concurrent requests using a semaphore.
+
+    This is separate from rate limiting - it controls how many
+    requests can be in-flight simultaneously.
+    """
+
+    def __init__(self, max_concurrency: int):
+        """
+        Initialize concurrency limiter.
+
+        Args:
+            max_concurrency: Maximum concurrent requests
+        """
+        self.max_concurrency = max_concurrency
+        self._semaphore = asyncio.Semaphore(max_concurrency)
+
+    async def __aenter__(self) -> "ConcurrencyLimiter":
+        await self._semaphore.acquire()
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        self._semaphore.release()
+
+    @property
+    def available_slots(self) -> int:
+        """Get number of available concurrency slots."""
+        # Semaphore._value is the internal counter
+        return self._semaphore._value  # noqa: SLF001