provide-foundation 0.0.0.dev0__py3-none-any.whl

provide/foundation/utils/rate_limiting.py

@@ -0,0 +1,112 @@
+"""
+Rate limiting utilities for Foundation.
+
+This module provides rate limiting implementations suitable for
+asynchronous applications, helping to manage request load and prevent abuse.
+"""
+
+import asyncio
+import time
+from typing import final
+
+
+@final
+class TokenBucketRateLimiter:
+    """
+    A Token Bucket rate limiter for asyncio applications.
+
+    This limiter allows for bursts up to a specified capacity and refills tokens
+    at a constant rate. It is designed to be thread-safe using an asyncio.Lock.
+    """
+
+    def __init__(self, capacity: float, refill_rate: float) -> None:
+        """
+        Initialize the TokenBucketRateLimiter.
+
+        Args:
+            capacity: The maximum number of tokens the bucket can hold
+                      (burst capacity).
+            refill_rate: The rate at which tokens are refilled per second.
+        """
+        if capacity <= 0:
+            raise ValueError("Capacity must be positive.")
+        if refill_rate <= 0:
+            raise ValueError("Refill rate must be positive.")
+
+        self._capacity: float = float(capacity)
+        self._refill_rate: float = float(refill_rate)
+        self._tokens: float = float(capacity)  # Start with a full bucket
+        self._last_refill_timestamp: float = time.monotonic()
+        self._lock: asyncio.Lock = asyncio.Lock()
+
+        # Cache logger instance to avoid repeated imports
+        self._logger = None
+        try:
+            from provide.foundation.logger import get_logger
+
+            self._logger = get_logger(__name__)
+            self._logger.debug(
+                "🔩🗑️ TokenBucketRateLimiter initialized: "
+                f"capacity={capacity}, refill_rate={refill_rate}"
+            )
+        except ImportError:
+            # Fallback if logger not available
+            pass
+
+    async def _refill_tokens(self) -> None:
+        """
+        Refills tokens based on the elapsed time since the last refill.
+        This method is not locked internally; caller must hold the lock.
+        """
+        now = time.monotonic()
+        elapsed_time = now - self._last_refill_timestamp
+        if elapsed_time > 0:  # only refill if time has passed
+            tokens_to_add = elapsed_time * self._refill_rate
+            # logger.debug(
+            #     f"🔩🗑️ Refilling: elapsed={elapsed_time:.4f}s, "
+            #     f"tokens_to_add={tokens_to_add:.4f}, "
+            #     f"current_tokens={self._tokens:.4f}"
+            # )
+            self._tokens = min(self._capacity, self._tokens + tokens_to_add)
+            self._last_refill_timestamp = now
+            # logger.debug(
+            #     f"🔩🗑️ Refilled: new_tokens={self._tokens:.4f}, "
+            #     f"last_refill_timestamp={self._last_refill_timestamp:.4f}"
+            # )
+
+    async def is_allowed(self) -> bool:
+        """
+        Check if a request is allowed based on available tokens.
+
+        This method is asynchronous and thread-safe. It refills tokens
+        based on elapsed time and then attempts to consume a token.
+
+        Returns:
+            True if the request is allowed, False otherwise.
+        """
+        async with self._lock:
+            await self._refill_tokens()  # Refill before checking
+
+            if self._tokens >= 1.0:
+                self._tokens -= 1.0
+                if self._logger:
+                    self._logger.debug(
+                        "🔩🗑️✅ Request allowed. Tokens remaining: "
+                        f"{self._tokens:.2f}/{self._capacity:.2f}"
+                    )
+                return True
+            else:
+                if self._logger:
+                    self._logger.warning(
+                        "🔩🗑️❌ Request denied. No tokens available. Tokens: "
+                        f"{self._tokens:.2f}/{self._capacity:.2f}"
+                    )
+                return False
+
+    async def get_current_tokens(self) -> float:
+        """Returns the current number of tokens, for testing/monitoring."""
+        async with self._lock:
+            # It might be useful to refill before getting, to get the most
+            # up-to-date count
+            # await self._refill_tokens()
+            return self._tokens

provide/foundation/utils/streams.py

@@ -0,0 +1,67 @@
+"""Stream utilities for foundation library."""
+
+import io
+import sys
+from typing import TextIO
+
+
+def get_safe_stderr() -> TextIO:
+    """Get a safe stderr stream, falling back to StringIO if stderr is not available.
+
+    This is used during initialization when sys.stderr might not be available
+    (e.g., in some embedded Python environments or during testing).
+
+    Returns:
+        A writable text stream, either sys.stderr or io.StringIO()
+    """
+    return (
+        sys.stderr
+        if hasattr(sys, "stderr") and sys.stderr is not None
+        else io.StringIO()
+    )
+
+
+def get_foundation_log_stream(output_setting: str) -> TextIO:
+    """Get the appropriate stream for Foundation internal logging.
+
+    Args:
+        output_setting: One of "stderr", "stdout", or "main"
+
+    Returns:
+        A writable text stream based on the output setting
+
+    Notes:
+        - "stderr": Returns sys.stderr (default, RPC-safe)
+        - "stdout": Returns sys.stdout
+        - "main": Returns the main logger stream from _PROVIDE_LOG_STREAM
+        - Invalid values default to sys.stderr with warning
+    """
+    if output_setting == "stdout":
+        return sys.stdout
+    elif output_setting == "main":
+        # Import here to avoid circular dependency
+        try:
+            from provide.foundation.streams import get_log_stream
+
+            return get_log_stream()
+        except ImportError:
+            # Fallback if setup module not available during initialization
+            return get_safe_stderr()
+    elif output_setting == "stderr":
+        return get_safe_stderr()
+    else:
+        # Invalid value - warn and default to stderr
+        # Import config logger here to avoid circular dependency
+        try:
+            from provide.foundation.logger.config import _get_config_logger
+
+            _get_config_logger().warning(
+                "[Foundation Config Warning] Invalid FOUNDATION_LOG_OUTPUT value, using stderr",
+                invalid_value=output_setting,
+                valid_options=["stderr", "stdout", "main"],
+                default_used="stderr",
+            )
+        except ImportError:
+            # During early initialization, just use stderr silently
+            pass
+        return get_safe_stderr()

provide/foundation/utils/timing.py

@@ -0,0 +1,93 @@
+"""
+Timing and performance utilities.
+
+Provides context managers and utilities for timing operations and logging performance.
+"""
+
+from collections.abc import Generator
+from contextlib import contextmanager
+import contextvars
+import time
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from provide.foundation.logger.base import FoundationLogger
+
+
+# Context variable for trace_id (applications can set this for request tracing)
+_PROVIDE_CONTEXT_TRACE_ID = contextvars.ContextVar(
+    "foundation_context_trace_id", default=None
+)
+
+
+@contextmanager
+def timed_block(
+    logger_instance: "FoundationLogger",
+    event_name: str,
+    layer_keys: dict[str, Any] | None = None,
+    initial_kvs: dict[str, Any] | None = None,
+    **extra_kvs: Any,
+) -> Generator[dict[str, Any], None, None]:
+    """
+    Context manager that logs the duration of a code block.
+
+    Logs at DEBUG when entering, INFO on success, ERROR on exception.
+
+    Args:
+        logger_instance: Logger to use for output
+        event_name: Name of the operation being timed
+        layer_keys: Semantic layer keys (e.g., llm-specific keys)
+        initial_kvs: Initial key-value pairs to include in logs
+        **extra_kvs: Additional key-value pairs
+
+    Yields:
+        A mutable dict that can be updated with additional context
+
+    Example:
+        >>> with timed_block(logger, "database_query") as ctx:
+        >>>     ctx["query"] = "SELECT * FROM users"
+        >>>     result = db.query("SELECT * FROM users")
+        >>>     ctx["rows"] = len(result)
+    """
+    # Combine all key-value pairs
+    all_kvs = {}
+    if layer_keys:
+        all_kvs.update(layer_keys)
+    if initial_kvs:
+        all_kvs.update(initial_kvs)
+    all_kvs.update(extra_kvs)
+
+    # Try to get trace_id from context
+    trace_id = _PROVIDE_CONTEXT_TRACE_ID.get()
+    if trace_id and "trace_id" not in all_kvs:
+        all_kvs["trace_id"] = trace_id
+
+    # Create context dict that can be modified
+    context: dict[str, Any] = {}
+
+    # Log start
+    logger_instance.debug(f"{event_name} started", **all_kvs)
+
+    start_time = time.perf_counter()
+    try:
+        yield context
+
+        # Success - calculate duration and log
+        duration = time.perf_counter() - start_time
+        all_kvs.update(context)
+        all_kvs["duration_seconds"] = round(duration, 3)
+        all_kvs["outcome"] = "success"
+
+        logger_instance.info(f"{event_name} completed", **all_kvs)
+
+    except Exception as e:
+        # Error - calculate duration and log with exception
+        duration = time.perf_counter() - start_time
+        all_kvs.update(context)
+        all_kvs["duration_seconds"] = round(duration, 3)
+        all_kvs["outcome"] = "error"
+        all_kvs["error.message"] = str(e)
+        all_kvs["error.type"] = type(e).__name__
+
+        logger_instance.error(f"{event_name} failed", exc_info=True, **all_kvs)
+        raise