flashlite 0.1.0__py3-none-any.whl

flashlite/observability/logging.py

@@ -0,0 +1,293 @@
+"""Structured logging for flashlite."""
+
+import json
+import logging
+import sys
+import time
+import uuid
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any, TextIO
+
+from ..types import CompletionRequest, CompletionResponse
+
+
+@dataclass
+class RequestLogEntry:
+    """A structured log entry for a completion request."""
+
+    request_id: str
+    timestamp: str
+    model: str
+    messages: list[dict[str, Any]]
+    parameters: dict[str, Any]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "type": "request",
+            "request_id": self.request_id,
+            "timestamp": self.timestamp,
+            "model": self.model,
+            "messages": self.messages,
+            "parameters": self.parameters,
+        }
+
+
+@dataclass
+class ResponseLogEntry:
+    """A structured log entry for a completion response."""
+
+    request_id: str
+    timestamp: str
+    model: str
+    content: str
+    finish_reason: str | None
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+    latency_ms: float
+    cached: bool = False
+    error: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "type": "response",
+            "request_id": self.request_id,
+            "timestamp": self.timestamp,
+            "model": self.model,
+            "content": self.content,
+            "finish_reason": self.finish_reason,
+            "usage": {
+                "input_tokens": self.input_tokens,
+                "output_tokens": self.output_tokens,
+                "total_tokens": self.total_tokens,
+            },
+            "latency_ms": self.latency_ms,
+            "cached": self.cached,
+            "error": self.error,
+        }
+
+
+class StructuredLogger:
+    """
+    A structured logger that outputs JSON-formatted log entries.
+
+    Can write to files, stdout, or both. Supports log rotation
+    and customizable formatting.
+
+    Example:
+        logger = StructuredLogger(
+            log_file="./logs/completions.jsonl",
+            log_level="INFO",
+            include_messages=True,
+        )
+
+        # Log a request
+        logger.log_request(request, request_id)
+
+        # Log a response
+        logger.log_response(response, request_id, latency_ms)
+    """
+
+    def __init__(
+        self,
+        log_file: str | Path | None = None,
+        log_level: str = "INFO",
+        include_messages: bool = True,
+        include_content: bool = True,
+        max_content_length: int | None = None,
+        redact_patterns: list[str] | None = None,
+        stdout: bool = False,
+    ):
+        """
+        Initialize the structured logger.
+
+        Args:
+            log_file: Path to log file (JSONL format). None disables file logging.
+            log_level: Minimum log level ("DEBUG", "INFO", "WARNING", "ERROR")
+            include_messages: Whether to include full message content in logs
+            include_content: Whether to include response content in logs
+            max_content_length: Max length of content to log (None = unlimited)
+            redact_patterns: Patterns to redact from logs (e.g., API keys)
+            stdout: Whether to also log to stdout
+        """
+        self._log_file: Path | None = Path(log_file) if log_file else None
+        self._log_level = getattr(logging, log_level.upper())
+        self._include_messages = include_messages
+        self._include_content = include_content
+        self._max_content_length = max_content_length
+        self._redact_patterns = redact_patterns or []
+        self._stdout = stdout
+        self._file_handle: TextIO | None = None
+
+        # Ensure log directory exists
+        if self._log_file:
+            self._log_file.parent.mkdir(parents=True, exist_ok=True)
+            self._file_handle = open(self._log_file, "a")
+
+    def _get_timestamp(self) -> str:
+        """Get current timestamp in ISO format."""
+        return datetime.now(UTC).isoformat()
+
+    def _redact(self, text: str) -> str:
+        """Redact sensitive patterns from text."""
+        for pattern in self._redact_patterns:
+            text = text.replace(pattern, "[REDACTED]")
+        return text
+
+    def _truncate(self, text: str) -> str:
+        """Truncate text if max length is set."""
+        if self._max_content_length and len(text) > self._max_content_length:
+            return text[: self._max_content_length] + "... [truncated]"
+        return text
+
+    def _write_entry(self, entry: dict[str, Any]) -> None:
+        """Write a log entry."""
+        json_str = json.dumps(entry, default=str)
+
+        if self._file_handle:
+            self._file_handle.write(json_str + "\n")
+            self._file_handle.flush()
+
+        if self._stdout:
+            print(json_str, file=sys.stdout)
+
+    def log_request(
+        self,
+        request: CompletionRequest,
+        request_id: str | None = None,
+    ) -> str:
+        """
+        Log a completion request.
+
+        Args:
+            request: The completion request
+            request_id: Optional request ID (generated if not provided)
+
+        Returns:
+            The request ID
+        """
+        if request_id is None:
+            request_id = str(uuid.uuid4())
+
+        # Build parameters dict
+        params: dict[str, Any] = {}
+        if request.temperature is not None:
+            params["temperature"] = request.temperature
+        if request.max_tokens is not None:
+            params["max_tokens"] = request.max_tokens
+        if request.max_completion_tokens is not None:
+            params["max_completion_tokens"] = request.max_completion_tokens
+        if request.top_p is not None:
+            params["top_p"] = request.top_p
+        if request.stop is not None:
+            params["stop"] = request.stop
+        if request.reasoning_effort is not None:
+            params["reasoning_effort"] = request.reasoning_effort
+        if request.thinking is not None:
+            params["thinking"] = request.thinking
+        params.update(request.extra_kwargs)
+
+        # Build messages
+        messages: list[dict[str, Any]] = []
+        if self._include_messages:
+            for msg in request.messages:
+                msg_dict = dict(msg)
+                if "content" in msg_dict:
+                    content = self._truncate(self._redact(str(msg_dict["content"])))
+                    msg_dict["content"] = content
+                messages.append(msg_dict)
+
+        entry = RequestLogEntry(
+            request_id=request_id,
+            timestamp=self._get_timestamp(),
+            model=request.model,
+            messages=messages,
+            parameters=params,
+        )
+
+        self._write_entry(entry.to_dict())
+        return request_id
+
+    def log_response(
+        self,
+        response: CompletionResponse,
+        request_id: str,
+        latency_ms: float,
+        cached: bool = False,
+    ) -> None:
+        """
+        Log a completion response.
+
+        Args:
+            response: The completion response
+            request_id: The corresponding request ID
+            latency_ms: Request latency in milliseconds
+            cached: Whether the response was from cache
+        """
+        content = ""
+        if self._include_content:
+            content = self._truncate(self._redact(response.content))
+
+        entry = ResponseLogEntry(
+            request_id=request_id,
+            timestamp=self._get_timestamp(),
+            model=response.model,
+            content=content,
+            finish_reason=response.finish_reason,
+            input_tokens=response.usage.input_tokens if response.usage else 0,
+            output_tokens=response.usage.output_tokens if response.usage else 0,
+            total_tokens=response.usage.total_tokens if response.usage else 0,
+            latency_ms=latency_ms,
+            cached=cached,
+        )
+
+        self._write_entry(entry.to_dict())
+
+    def log_error(
+        self,
+        request_id: str,
+        error: Exception,
+        latency_ms: float,
+    ) -> None:
+        """
+        Log an error response.
+
+        Args:
+            request_id: The corresponding request ID
+            error: The exception that occurred
+            latency_ms: Request latency in milliseconds
+        """
+        entry = {
+            "type": "error",
+            "request_id": request_id,
+            "timestamp": self._get_timestamp(),
+            "error": str(error),
+            "error_type": type(error).__name__,
+            "latency_ms": latency_ms,
+        }
+
+        self._write_entry(entry)
+
+    def close(self) -> None:
+        """Close the log file."""
+        if self._file_handle:
+            self._file_handle.close()
+            self._file_handle = None
+
+
+@dataclass
+class RequestContext:
+    """Context for tracking a single request through the pipeline."""
+
+    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    start_time: float = field(default_factory=time.perf_counter)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def elapsed_ms(self) -> float:
+        """Get elapsed time in milliseconds."""
+        return (time.perf_counter() - self.start_time) * 1000

flashlite/observability/metrics.py

@@ -0,0 +1,221 @@
+"""Metrics and cost tracking for flashlite."""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..types import CompletionResponse
+
+logger = logging.getLogger(__name__)
+
+
+# Approximate costs per 1K tokens (USD) as of early 2025
+# These are estimates and may be outdated - use litellm's cost tracking for accuracy
+DEFAULT_COSTS: dict[str, dict[str, float]] = {
+    # OpenAI
+    "gpt-4o": {"input": 0.0025, "output": 0.01},
+    "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
+    "gpt-4-turbo": {"input": 0.01, "output": 0.03},
+    "gpt-4": {"input": 0.03, "output": 0.06},
+    "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
+    "o1": {"input": 0.015, "output": 0.06},
+    "o1-mini": {"input": 0.003, "output": 0.012},
+    "o3-mini": {"input": 0.003, "output": 0.012},
+    # Anthropic
+    "claude-3-5-sonnet-20241022": {"input": 0.003, "output": 0.015},
+    "claude-sonnet-4-20250514": {"input": 0.003, "output": 0.015},
+    "claude-3-5-haiku-20241022": {"input": 0.001, "output": 0.005},
+    "claude-3-opus-20240229": {"input": 0.015, "output": 0.075},
+    # Add more as needed
+}
+
+
+@dataclass
+class CostMetrics:
+    """Accumulated cost metrics."""
+
+    total_input_tokens: int = 0
+    total_output_tokens: int = 0
+    total_requests: int = 0
+    total_cost_usd: float = 0.0
+    cost_by_model: dict[str, float] = field(default_factory=dict)
+    tokens_by_model: dict[str, dict[str, int]] = field(default_factory=dict)
+
+
+class CostTracker:
+    """
+    Tracks token usage and estimated costs across requests.
+
+    Example:
+        tracker = CostTracker(budget_limit=10.0)
+
+        # Track a response
+        tracker.track(response)
+
+        # Check totals
+        print(f"Total cost: ${tracker.total_cost:.2f}")
+        print(f"Budget remaining: ${tracker.budget_remaining:.2f}")
+
+        # Export report
+        report = tracker.get_report()
+    """
+
+    def __init__(
+        self,
+        budget_limit: float | None = None,
+        warn_at_percent: float = 80.0,
+        custom_costs: dict[str, dict[str, float]] | None = None,
+    ):
+        """
+        Initialize the cost tracker.
+
+        Args:
+            budget_limit: Maximum budget in USD (None = no limit)
+            warn_at_percent: Warn when this percentage of budget is used
+            custom_costs: Custom cost overrides per model
+        """
+        self._budget_limit = budget_limit
+        self._warn_at_percent = warn_at_percent
+        self._costs = {**DEFAULT_COSTS, **(custom_costs or {})}
+        self._metrics = CostMetrics()
+        self._budget_warning_issued = False
+
+    def _get_model_cost(self, model: str) -> dict[str, float]:
+        """Get cost per 1K tokens for a model."""
+        # Exact match
+        if model in self._costs:
+            return self._costs[model]
+
+        # Try to match by prefix (e.g., "gpt-4o" matches "gpt-4o-2024-...")
+        model_lower = model.lower()
+        for known_model, cost in self._costs.items():
+            if model_lower.startswith(known_model.lower()):
+                return cost
+
+        # Default to GPT-4o pricing as a reasonable estimate
+        logger.debug(f"Unknown model cost for '{model}', using gpt-4o pricing estimate")
+        return self._costs.get("gpt-4o", {"input": 0.0025, "output": 0.01})
+
+    def track(self, response: CompletionResponse) -> float:
+        """
+        Track a completion response and return its cost.
+
+        Args:
+            response: The completion response to track
+
+        Returns:
+            The estimated cost in USD for this response
+        """
+        if not response.usage:
+            return 0.0
+
+        input_tokens = response.usage.input_tokens
+        output_tokens = response.usage.output_tokens
+        model = response.model
+
+        # Calculate cost
+        model_costs = self._get_model_cost(model)
+        input_cost = (input_tokens / 1000) * model_costs["input"]
+        output_cost = (output_tokens / 1000) * model_costs["output"]
+        total_cost = input_cost + output_cost
+
+        # Update metrics
+        self._metrics.total_input_tokens += input_tokens
+        self._metrics.total_output_tokens += output_tokens
+        self._metrics.total_requests += 1
+        self._metrics.total_cost_usd += total_cost
+
+        # Per-model tracking
+        if model not in self._metrics.cost_by_model:
+            self._metrics.cost_by_model[model] = 0.0
+            self._metrics.tokens_by_model[model] = {"input": 0, "output": 0}
+
+        self._metrics.cost_by_model[model] += total_cost
+        self._metrics.tokens_by_model[model]["input"] += input_tokens
+        self._metrics.tokens_by_model[model]["output"] += output_tokens
+
+        # Check budget
+        self._check_budget()
+
+        return total_cost
+
+    def _check_budget(self) -> None:
+        """Check if budget thresholds are exceeded."""
+        if self._budget_limit is None:
+            return
+
+        percent_used = (self._metrics.total_cost_usd / self._budget_limit) * 100
+
+        # Warning threshold
+        if percent_used >= self._warn_at_percent and not self._budget_warning_issued:
+            logger.warning(
+                f"Budget warning: {percent_used:.1f}% of ${self._budget_limit:.2f} budget used "
+                f"(${self._metrics.total_cost_usd:.4f} spent)"
+            )
+            self._budget_warning_issued = True
+
+        # Hard limit
+        if self._metrics.total_cost_usd >= self._budget_limit:
+            raise BudgetExceededError(
+                f"Budget limit of ${self._budget_limit:.2f} exceeded "
+                f"(${self._metrics.total_cost_usd:.4f} spent)"
+            )
+
+    @property
+    def total_cost(self) -> float:
+        """Total cost in USD."""
+        return self._metrics.total_cost_usd
+
+    @property
+    def total_tokens(self) -> int:
+        """Total tokens used."""
+        return self._metrics.total_input_tokens + self._metrics.total_output_tokens
+
+    @property
+    def total_requests(self) -> int:
+        """Total number of requests tracked."""
+        return self._metrics.total_requests
+
+    @property
+    def budget_remaining(self) -> float | None:
+        """Remaining budget in USD, or None if no limit."""
+        if self._budget_limit is None:
+            return None
+        return max(0.0, self._budget_limit - self._metrics.total_cost_usd)
+
+    def get_report(self) -> dict[str, Any]:
+        """
+        Get a detailed cost report.
+
+        Returns:
+            Dictionary with cost breakdown
+        """
+        return {
+            "total_cost_usd": self._metrics.total_cost_usd,
+            "total_requests": self._metrics.total_requests,
+            "total_tokens": {
+                "input": self._metrics.total_input_tokens,
+                "output": self._metrics.total_output_tokens,
+                "total": self.total_tokens,
+            },
+            "budget_limit_usd": self._budget_limit,
+            "budget_remaining_usd": self.budget_remaining,
+            "by_model": {
+                model: {
+                    "cost_usd": self._metrics.cost_by_model[model],
+                    "tokens": self._metrics.tokens_by_model[model],
+                }
+                for model in self._metrics.cost_by_model
+            },
+        }
+
+    def reset(self) -> None:
+        """Reset all metrics."""
+        self._metrics = CostMetrics()
+        self._budget_warning_issued = False
+
+
+class BudgetExceededError(Exception):
+    """Raised when the budget limit is exceeded."""
+
+    pass

flashlite/structured/__init__.py

@@ -0,0 +1,31 @@
+"""Structured outputs module for Pydantic model integration."""
+
+from .outputs import (
+    StructuredOutputError,
+    extract_json_from_content,
+    format_validation_error_for_retry,
+    parse_json_response,
+    validate_response,
+)
+from .schema import (
+    format_schema_for_openai,
+    generate_json_schema,
+    get_field_descriptions,
+    is_supported_type,
+    schema_to_prompt,
+)
+
+__all__ = [
+    # Schema generation
+    "generate_json_schema",
+    "schema_to_prompt",
+    "get_field_descriptions",
+    "format_schema_for_openai",
+    "is_supported_type",
+    # Parsing and validation
+    "parse_json_response",
+    "validate_response",
+    "format_validation_error_for_retry",
+    "extract_json_from_content",
+    "StructuredOutputError",
+]