flashlite 0.1.0__py3-none-any.whl

flashlite/middleware/retry.py

@@ -0,0 +1,149 @@
+"""Retry middleware with exponential backoff."""
+
+import asyncio
+import logging
+import random
+
+from tenacity import (
+    AsyncRetrying,
+    RetryError,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential_jitter,
+)
+
+from ..types import (
+    CompletionError,
+    CompletionRequest,
+    CompletionResponse,
+    RetryConfig,
+)
+from .base import CompletionHandler, Middleware
+
+logger = logging.getLogger(__name__)
+
+
+def _should_retry(exception: BaseException) -> bool:
+    """Determine if an exception should trigger a retry."""
+    if isinstance(exception, CompletionError):
+        # Retry on specific status codes
+        if exception.status_code in (429, 500, 502, 503, 504):
+            return True
+        # Don't retry on client errors (4xx except 429)
+        if exception.status_code and 400 <= exception.status_code < 500:
+            return False
+    # Retry on connection errors, timeouts, etc.
+    if isinstance(exception, (ConnectionError, TimeoutError, asyncio.TimeoutError)):
+        return True
+    return False
+
+
+class RetryMiddleware(Middleware):
+    """
+    Middleware that retries failed requests with exponential backoff.
+
+    Uses tenacity for retry logic with:
+    - Exponential backoff with jitter
+    - Configurable max attempts and delays
+    - Retries on transient errors (429, 5xx, connection errors)
+    """
+
+    def __init__(self, config: RetryConfig | None = None):
+        """
+        Initialize retry middleware.
+
+        Args:
+            config: Retry configuration. Uses defaults if not provided.
+        """
+        self.config = config or RetryConfig()
+
+    async def __call__(
+        self,
+        request: CompletionRequest,
+        next_handler: CompletionHandler,
+    ) -> CompletionResponse:
+        """Execute with retry logic."""
+        attempt = 0
+
+        try:
+            async for attempt_state in AsyncRetrying(
+                stop=stop_after_attempt(self.config.max_attempts),
+                wait=wait_exponential_jitter(
+                    initial=self.config.initial_delay,
+                    max=self.config.max_delay,
+                    exp_base=self.config.exponential_base,
+                    jitter=self.config.initial_delay if self.config.jitter else 0,
+                ),
+                retry=retry_if_exception(_should_retry),
+                reraise=True,
+            ):
+                with attempt_state:
+                    attempt = attempt_state.retry_state.attempt_number
+                    if attempt > 1:
+                        logger.info(
+                            f"Retry attempt {attempt}/{self.config.max_attempts} "
+                            f"for model={request.model}"
+                        )
+                    return await next_handler(request)
+
+        except RetryError as e:
+            # Re-raise the last exception
+            if e.last_attempt.failed:
+                exc = e.last_attempt.exception()
+                if exc is not None:
+                    raise exc from e
+            raise CompletionError("Retry attempts exhausted") from e
+
+        # This should never be reached, but satisfies type checker
+        raise CompletionError("Retry logic failed unexpectedly")
+
+
+class SimpleRetryMiddleware(Middleware):
+    """
+    A simpler retry implementation without tenacity dependency.
+
+    Useful for understanding the retry logic or when tenacity isn't available.
+    """
+
+    def __init__(self, config: RetryConfig | None = None):
+        self.config = config or RetryConfig()
+
+    async def __call__(
+        self,
+        request: CompletionRequest,
+        next_handler: CompletionHandler,
+    ) -> CompletionResponse:
+        last_exception: Exception | None = None
+        delay = self.config.initial_delay
+
+        for attempt in range(1, self.config.max_attempts + 1):
+            try:
+                return await next_handler(request)
+            except Exception as e:
+                last_exception = e
+
+                if not _should_retry(e):
+                    raise
+
+                if attempt == self.config.max_attempts:
+                    raise
+
+                # Calculate delay with optional jitter
+                actual_delay = delay
+                if self.config.jitter:
+                    actual_delay = delay * (0.5 + random.random())
+
+                logger.info(
+                    f"Attempt {attempt}/{self.config.max_attempts} failed, "
+                    f"retrying in {actual_delay:.2f}s: {e}"
+                )
+
+                await asyncio.sleep(actual_delay)
+
+                # Exponential backoff for next attempt
+                delay = min(delay * self.config.exponential_base, self.config.max_delay)
+
+        # Should never reach here
+        if last_exception:
+            raise last_exception
+        raise CompletionError("Retry logic failed unexpectedly")

flashlite/observability/__init__.py

@@ -0,0 +1,34 @@
+"""Observability module for flashlite."""
+
+from .callbacks import (
+    CallbackManager,
+    OnErrorCallback,
+    OnRequestCallback,
+    OnResponseCallback,
+    create_logging_callbacks,
+)
+from .inspect_compat import FlashliteModelAPI, InspectLogEntry, InspectLogger
+from .logging import RequestContext, RequestLogEntry, ResponseLogEntry, StructuredLogger
+from .metrics import BudgetExceededError, CostMetrics, CostTracker
+
+__all__ = [
+    # Logging
+    "StructuredLogger",
+    "RequestLogEntry",
+    "ResponseLogEntry",
+    "RequestContext",
+    # Metrics
+    "CostTracker",
+    "CostMetrics",
+    "BudgetExceededError",
+    # Callbacks
+    "CallbackManager",
+    "OnRequestCallback",
+    "OnResponseCallback",
+    "OnErrorCallback",
+    "create_logging_callbacks",
+    # Inspect
+    "InspectLogger",
+    "InspectLogEntry",
+    "FlashliteModelAPI",
+]

flashlite/observability/callbacks.py

@@ -0,0 +1,155 @@
+"""Callback system for flashlite observability."""
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..types import CompletionRequest, CompletionResponse
+
+# Callback type definitions
+OnRequestCallback = Callable[[CompletionRequest, str], Awaitable[None] | None]
+OnResponseCallback = Callable[
+    [CompletionResponse, str, float, bool], Awaitable[None] | None
+]
+OnErrorCallback = Callable[[Exception, str, float], Awaitable[None] | None]
+
+
+@dataclass
+class CallbackManager:
+    """
+    Manages callbacks for request/response lifecycle events.
+
+    Example:
+        callbacks = CallbackManager()
+
+        @callbacks.on_request
+        async def log_request(request, request_id):
+            print(f"Request {request_id}: {request.model}")
+
+        @callbacks.on_response
+        async def log_response(response, request_id, latency_ms, cached):
+            print(f"Response {request_id}: {latency_ms}ms")
+
+        # Or register directly
+        callbacks.add_on_request(my_callback)
+    """
+
+    _on_request: list[OnRequestCallback] = field(default_factory=list)
+    _on_response: list[OnResponseCallback] = field(default_factory=list)
+    _on_error: list[OnErrorCallback] = field(default_factory=list)
+
+    def add_on_request(self, callback: OnRequestCallback) -> None:
+        """Add a callback to be called before each request."""
+        self._on_request.append(callback)
+
+    def add_on_response(self, callback: OnResponseCallback) -> None:
+        """Add a callback to be called after each successful response."""
+        self._on_response.append(callback)
+
+    def add_on_error(self, callback: OnErrorCallback) -> None:
+        """Add a callback to be called on request errors."""
+        self._on_error.append(callback)
+
+    def on_request(
+        self, callback: OnRequestCallback
+    ) -> OnRequestCallback:
+        """Decorator to register a request callback."""
+        self.add_on_request(callback)
+        return callback
+
+    def on_response(
+        self, callback: OnResponseCallback
+    ) -> OnResponseCallback:
+        """Decorator to register a response callback."""
+        self.add_on_response(callback)
+        return callback
+
+    def on_error(self, callback: OnErrorCallback) -> OnErrorCallback:
+        """Decorator to register an error callback."""
+        self.add_on_error(callback)
+        return callback
+
+    async def emit_request(
+        self,
+        request: CompletionRequest,
+        request_id: str,
+    ) -> None:
+        """Emit a request event to all registered callbacks."""
+        for callback in self._on_request:
+            result = callback(request, request_id)
+            if isinstance(result, Awaitable):
+                await result
+
+    async def emit_response(
+        self,
+        response: CompletionResponse,
+        request_id: str,
+        latency_ms: float,
+        cached: bool = False,
+    ) -> None:
+        """Emit a response event to all registered callbacks."""
+        for callback in self._on_response:
+            result = callback(response, request_id, latency_ms, cached)
+            if isinstance(result, Awaitable):
+                await result
+
+    async def emit_error(
+        self,
+        error: Exception,
+        request_id: str,
+        latency_ms: float,
+    ) -> None:
+        """Emit an error event to all registered callbacks."""
+        for callback in self._on_error:
+            result = callback(error, request_id, latency_ms)
+            if isinstance(result, Awaitable):
+                await result
+
+
+def create_logging_callbacks(
+    logger: Any,
+    level: str = "INFO",
+) -> CallbackManager:
+    """
+    Create a CallbackManager with standard logging callbacks.
+
+    Args:
+        logger: A logging.Logger instance
+        level: Log level to use
+
+    Returns:
+        A configured CallbackManager
+    """
+    import logging as stdlib_logging
+
+    log_level = getattr(stdlib_logging, level.upper())
+    callbacks = CallbackManager()
+
+    @callbacks.on_request
+    def log_request(request: CompletionRequest, request_id: str) -> None:
+        logger.log(
+            log_level,
+            f"[{request_id[:8]}] Request: model={request.model}",
+        )
+
+    @callbacks.on_response
+    def log_response(
+        response: CompletionResponse,
+        request_id: str,
+        latency_ms: float,
+        cached: bool,
+    ) -> None:
+        cache_str = " (cached)" if cached else ""
+        tokens = response.usage.total_tokens if response.usage else 0
+        logger.log(
+            log_level,
+            f"[{request_id[:8]}] Response: {latency_ms:.1f}ms, {tokens} tokens{cache_str}",
+        )
+
+    @callbacks.on_error
+    def log_error(error: Exception, request_id: str, latency_ms: float) -> None:
+        logger.error(
+            f"[{request_id[:8]}] Error after {latency_ms:.1f}ms: {error}",
+        )
+
+    return callbacks

flashlite/observability/inspect_compat.py

@@ -0,0 +1,266 @@
+"""Inspect framework compatibility layer for flashlite.
+
+This module provides interoperability with the UK AISI's Inspect framework
+(https://inspect.ai-safety-institute.org.uk/).
+
+It includes:
+- Log format compatible with Inspect's eval logging
+- ModelAPI protocol implementation for use as an Inspect solver backend
+- Hooks for Inspect's TaskState integration
+"""
+
+import json
+import logging
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from ..types import CompletionRequest, CompletionResponse
+
+if TYPE_CHECKING:
+    from ..client import Flashlite
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class InspectLogEntry:
+    """A log entry in Inspect-compatible format."""
+
+    eval_id: str
+    sample_id: str | int
+    epoch: int
+    model: str
+    input: list[dict[str, Any]]
+    output: str
+    tokens: dict[str, int]
+    timestamp: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to Inspect log format."""
+        return {
+            "eval_id": self.eval_id,
+            "sample_id": self.sample_id,
+            "epoch": self.epoch,
+            "model": self.model,
+            "input": self.input,
+            "output": self.output,
+            "tokens": self.tokens,
+            "timestamp": self.timestamp,
+            "metadata": self.metadata,
+        }
+
+
+class InspectLogger:
+    """
+    A logger that outputs in Inspect-compatible format.
+
+    This allows Flashlite logs to be analyzed alongside Inspect eval logs,
+    enabling unified observability across evaluation runs.
+
+    Example:
+        inspect_logger = InspectLogger(
+            log_dir="./logs",
+            eval_id="my-eval-001",
+        )
+
+        # Log a completion
+        inspect_logger.log(
+            request=request,
+            response=response,
+            sample_id="sample_123",
+            epoch=0,
+        )
+
+        # Close when done
+        inspect_logger.close()
+    """
+
+    def __init__(
+        self,
+        log_dir: str | Path,
+        eval_id: str | None = None,
+        append: bool = True,
+    ):
+        """
+        Initialize the Inspect logger.
+
+        Args:
+            log_dir: Directory to write log files
+            eval_id: Evaluation ID (auto-generated if not provided)
+            append: Whether to append to existing log file
+        """
+        self._log_dir = Path(log_dir)
+        self._log_dir.mkdir(parents=True, exist_ok=True)
+
+        self._eval_id = eval_id or self._generate_eval_id()
+        self._log_file = self._log_dir / f"{self._eval_id}.jsonl"
+        self._mode = "a" if append else "w"
+        self._file_handle = open(self._log_file, self._mode)
+        self._sample_count = 0
+
+    def _generate_eval_id(self) -> str:
+        """Generate a unique evaluation ID."""
+        timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
+        return f"flashlite_eval_{timestamp}"
+
+    def log(
+        self,
+        request: CompletionRequest,
+        response: CompletionResponse,
+        sample_id: str | int | None = None,
+        epoch: int = 0,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Log a request/response pair in Inspect format.
+
+        Args:
+            request: The completion request
+            response: The completion response
+            sample_id: Sample identifier (auto-incremented if not provided)
+            epoch: Epoch number for multi-epoch evals
+            metadata: Additional metadata to include
+        """
+        if sample_id is None:
+            sample_id = self._sample_count
+            self._sample_count += 1
+
+        # Convert messages to Inspect format
+        input_messages = [
+            {"role": msg.get("role", "user"), "content": msg.get("content", "")}
+            for msg in request.messages
+        ]
+
+        entry = InspectLogEntry(
+            eval_id=self._eval_id,
+            sample_id=sample_id,
+            epoch=epoch,
+            model=response.model,
+            input=input_messages,
+            output=response.content,
+            tokens={
+                "input": response.usage.input_tokens if response.usage else 0,
+                "output": response.usage.output_tokens if response.usage else 0,
+                "total": response.usage.total_tokens if response.usage else 0,
+            },
+            timestamp=datetime.now(UTC).isoformat(),
+            metadata=metadata or {},
+        )
+
+        json_str = json.dumps(entry.to_dict())
+        self._file_handle.write(json_str + "\n")
+        self._file_handle.flush()
+
+    def close(self) -> None:
+        """Close the log file."""
+        if self._file_handle:
+            self._file_handle.close()
+
+    @property
+    def eval_id(self) -> str:
+        """Get the evaluation ID."""
+        return self._eval_id
+
+    @property
+    def log_file(self) -> Path:
+        """Get the log file path."""
+        return self._log_file
+
+
+class FlashliteModelAPI:
+    """
+    An adapter that implements a ModelAPI-like interface for Inspect integration.
+
+    This allows Flashlite to be used as a model backend in Inspect evaluations.
+
+    Example:
+        from flashlite import Flashlite
+        from flashlite.observability import FlashliteModelAPI
+
+        # Create Flashlite client
+        client = Flashlite(rate_limit=RateLimitConfig(requests_per_minute=60))
+
+        # Wrap for Inspect
+        model_api = FlashliteModelAPI(client, model="gpt-4o")
+
+        # Use in Inspect eval (pseudocode)
+        # @task
+        # def my_eval():
+        #     return Task(
+        #         dataset=my_dataset,
+        #         solver=my_solver,
+        #         model=model_api,
+        #     )
+    """
+
+    def __init__(
+        self,
+        client: "Flashlite",
+        model: str | None = None,
+        **default_kwargs: Any,
+    ):
+        """
+        Initialize the Inspect model adapter.
+
+        Args:
+            client: The Flashlite client to use
+            model: Default model to use (can be overridden per-request)
+            **default_kwargs: Default parameters for completions
+        """
+        self._client = client
+        self._model = model
+        self._default_kwargs = default_kwargs
+
+    async def generate(
+        self,
+        messages: list[dict[str, Any]],
+        model: str | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """
+        Generate a completion (Inspect-compatible interface).
+
+        Args:
+            messages: List of messages
+            model: Model to use (overrides default)
+            **kwargs: Additional parameters
+
+        Returns:
+            Inspect-compatible response dict
+        """
+        # Merge kwargs
+        call_kwargs = {**self._default_kwargs, **kwargs}
+
+        # Call Flashlite
+        response = await self._client.complete(
+            model=model or self._model,
+            messages=messages,
+            **call_kwargs,
+        )
+
+        # Return in Inspect-compatible format
+        return {
+            "choices": [
+                {
+                    "message": {
+                        "role": "assistant",
+                        "content": response.content,
+                    },
+                    "finish_reason": response.finish_reason,
+                }
+            ],
+            "usage": {
+                "prompt_tokens": response.usage.input_tokens if response.usage else 0,
+                "completion_tokens": response.usage.output_tokens if response.usage else 0,
+                "total_tokens": response.usage.total_tokens if response.usage else 0,
+            },
+            "model": response.model,
+        }
+
+    @property
+    def model_name(self) -> str | None:
+        """Get the default model name."""
+        return self._model