krons 0.1.0__py3-none-any.whl

kronos/services/utilities/header_factory.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""HTTP header construction utilities for API authentication."""
+
+from typing import Literal
+
+from pydantic import SecretStr
+
+AUTH_TYPES = Literal["bearer", "x-api-key", "none"]
+
+
+class HeaderFactory:
+    """Factory for constructing HTTP headers with various auth schemes.
+
+    Supports Bearer token, x-api-key, and no-auth patterns.
+    Handles SecretStr unwrapping and validation automatically.
+
+    Example:
+        >>> headers = HeaderFactory.get_header("bearer", api_key="sk-xxx")
+        >>> headers
+        {'Content-Type': 'application/json', 'Authorization': 'Bearer sk-xxx'}
+    """
+
+    @staticmethod
+    def get_content_type_header(
+        content_type: str = "application/json",
+    ) -> dict[str, str]:
+        """Build Content-Type header dict."""
+        return {"Content-Type": content_type}
+
+    @staticmethod
+    def get_bearer_auth_header(api_key: str) -> dict[str, str]:
+        """Build Authorization header with Bearer scheme."""
+        return {"Authorization": f"Bearer {api_key}"}
+
+    @staticmethod
+    def get_x_api_key_header(api_key: str) -> dict[str, str]:
+        """Build x-api-key header for providers requiring this scheme."""
+        return {"x-api-key": api_key}
+
+    @staticmethod
+    def get_header(
+        auth_type: AUTH_TYPES,
+        content_type: str | None = "application/json",
+        api_key: str | SecretStr | None = None,
+        default_headers: dict[str, str] | None = None,
+    ) -> dict[str, str]:
+        """Construct complete HTTP headers for API requests.
+
+        Args:
+            auth_type: Authentication scheme ("bearer", "x-api-key", "none").
+            content_type: Content-Type value (None to omit).
+            api_key: API key (str or SecretStr, required unless auth_type="none").
+            default_headers: Additional headers to merge.
+
+        Returns:
+            Complete header dict ready for HTTP client.
+
+        Raises:
+            ValueError: If api_key missing/empty when auth required, or invalid auth_type.
+        """
+        dict_ = {}
+        if content_type is not None:
+            dict_ = HeaderFactory.get_content_type_header(content_type)
+
+        if auth_type == "none":
+            pass
+        else:
+            if isinstance(api_key, SecretStr):
+                api_key = api_key.get_secret_value()
+
+            if not api_key or not str(api_key).strip():
+                raise ValueError("API key is required for authentication")
+
+            api_key = api_key.strip()
+
+            if auth_type == "bearer":
+                dict_.update(HeaderFactory.get_bearer_auth_header(api_key))
+            elif auth_type == "x-api-key":
+                dict_.update(HeaderFactory.get_x_api_key_header(api_key))
+            else:
+                raise ValueError(f"Unsupported auth type: {auth_type}")
+
+        if default_headers:
+            dict_.update(default_headers)
+        return dict_

kronos/services/utilities/rate_limited_executor.py

@@ -0,0 +1,271 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Rate-limited execution infrastructure with dual token bucket support.
+
+Provides permission-based rate limiting for API calls with separate
+request count and token usage limits, plus atomic rollback on partial acquire.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self
+
+from typing_extensions import override
+
+from kronos.core import Event, Executor, Processor
+from kronos.services.endpoint import APICalling
+from kronos.utils.concurrency import get_cancelled_exc_class, sleep
+
+from .rate_limiter import TokenBucket
+
+if TYPE_CHECKING:
+    import asyncio
+
+    from kronos.core import Pile
+
+__all__ = ("RateLimitedExecutor", "RateLimitedProcessor")
+
+logger = logging.getLogger(__name__)
+
+
+class RateLimitedProcessor(Processor):
+    """Processor with dual token bucket rate limiting (requests + tokens).
+
+    Enforces both request count and token usage limits atomically.
+    Automatically rolls back request bucket if token bucket acquire fails.
+
+    Example:
+        >>> req_bucket = TokenBucket(RateLimitConfig(capacity=100, refill_rate=1.67))
+        >>> tok_bucket = TokenBucket(RateLimitConfig(capacity=100000, refill_rate=1667))
+        >>> processor = await RateLimitedProcessor.create(
+        ...     queue_capacity=50, capacity_refresh_time=60.0,
+        ...     request_bucket=req_bucket, token_bucket=tok_bucket
+        ... )
+    """
+
+    event_type = APICalling
+
+    def __init__(
+        self,
+        queue_capacity: int,
+        capacity_refresh_time: float,
+        pile: Pile[Event] | None = None,
+        executor: Executor | None = None,
+        request_bucket: TokenBucket | None = None,
+        token_bucket: TokenBucket | None = None,
+        replenishment_interval: float = 60.0,
+        concurrency_limit: int = 100,
+        max_queue_size: int = 1000,
+        max_denial_tracking: int = 10000,
+    ) -> None:
+        """Initialize rate-limited processor.
+
+        Args:
+            queue_capacity: Max events per batch.
+            capacity_refresh_time: Batch refresh interval (seconds).
+            pile: Reference to executor's Flow.items (set by executor).
+            executor: Reference to executor for progression updates.
+            request_bucket: TokenBucket for request rate limiting.
+            token_bucket: TokenBucket for token rate limiting.
+            replenishment_interval: Rate limit reset interval.
+            concurrency_limit: Max concurrent executions.
+            max_queue_size: Max queue size.
+            max_denial_tracking: Max denial entries to track.
+        """
+        super().__init__(  # type: ignore[arg-type]
+            queue_capacity=queue_capacity,
+            capacity_refresh_time=capacity_refresh_time,
+            pile=pile,  # type: ignore[arg-type]
+            executor=executor,
+            concurrency_limit=concurrency_limit,
+            max_queue_size=max_queue_size,
+            max_denial_tracking=max_denial_tracking,
+        )
+
+        self.request_bucket = request_bucket
+        self.token_bucket = token_bucket
+        self.replenishment_interval = replenishment_interval
+        self.concurrency_limit = concurrency_limit
+        self._replenisher_task: asyncio.Task[None] | None = None
+
+    async def start_replenishing(self) -> None:
+        """Background task: periodically reset rate limit buckets to full capacity."""
+        await self.start()
+
+        try:
+            while not self.is_stopped():
+                await sleep(self.replenishment_interval)
+
+                if self.request_bucket:
+                    await self.request_bucket.reset()
+                    logger.debug(
+                        "Request bucket replenished: %d requests",
+                        self.request_bucket.capacity,
+                    )
+
+                if self.token_bucket:
+                    await self.token_bucket.reset()
+                    logger.debug(
+                        "Token bucket replenished: %d tokens",
+                        self.token_bucket.capacity,
+                    )
+
+        except get_cancelled_exc_class():
+            logger.info("Rate limit replenisher task cancelled.")
+
+    @override
+    @classmethod
+    async def create(  # type: ignore[override]
+        cls,
+        queue_capacity: int,
+        capacity_refresh_time: float,
+        pile: Pile[Event] | None = None,
+        executor: Executor | None = None,
+        request_bucket: TokenBucket | None = None,
+        token_bucket: TokenBucket | None = None,
+        replenishment_interval: float = 60.0,
+        concurrency_limit: int = 100,
+        max_queue_size: int = 1000,
+        max_denial_tracking: int = 10000,
+    ) -> Self:
+        """Factory: create processor and start background replenishment task."""
+        self = cls(
+            queue_capacity=queue_capacity,
+            capacity_refresh_time=capacity_refresh_time,
+            pile=pile,
+            executor=executor,
+            request_bucket=request_bucket,
+            token_bucket=token_bucket,
+            replenishment_interval=replenishment_interval,
+            concurrency_limit=concurrency_limit,
+            max_queue_size=max_queue_size,
+            max_denial_tracking=max_denial_tracking,
+        )
+
+        import asyncio
+
+        self._replenisher_task = asyncio.create_task(self.start_replenishing())
+
+        return self
+
+    @override
+    async def stop(self) -> None:
+        """Stop processor and cancel background replenishment task."""
+        if self._replenisher_task:
+            self._replenisher_task.cancel()
+            try:
+                await self._replenisher_task
+            except get_cancelled_exc_class():
+                pass
+
+        await super().stop()
+
+    @override
+    async def request_permission(
+        self,
+        required_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> bool:
+        """Check rate limits and acquire tokens atomically.
+
+        Acquires from request bucket first, then token bucket. If token bucket
+        fails, rolls back request bucket automatically.
+
+        Args:
+            required_tokens: Token count for this request (None = skip token check).
+            **kwargs: Ignored (for interface compatibility).
+
+        Returns:
+            True if permitted, False if rate limited.
+        """
+        if self.request_bucket is None and self.token_bucket is None:
+            return True
+
+        request_acquired = False
+        if self.request_bucket:
+            request_acquired = await self.request_bucket.try_acquire(tokens=1)
+            if not request_acquired:
+                logger.debug("Request rate limit exceeded")
+                return False
+
+        if self.token_bucket and required_tokens:
+            token_acquired = await self.token_bucket.try_acquire(tokens=required_tokens)
+            if not token_acquired:
+                if request_acquired and self.request_bucket:
+                    await self.request_bucket.release(tokens=1)
+
+                logger.debug(
+                    f"Token rate limit exceeded (required: {required_tokens}, "
+                    f"available: {self.token_bucket.tokens:.0f})"
+                )
+                return False
+
+        return True
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize processor config to dict (excludes runtime state)."""
+        return {
+            "queue_capacity": self.queue_capacity,
+            "capacity_refresh_time": self.capacity_refresh_time,
+            "replenishment_interval": self.replenishment_interval,
+            "concurrency_limit": self.concurrency_limit,
+            "max_queue_size": self.max_queue_size,
+            "max_denial_tracking": self.max_denial_tracking,
+            "request_bucket": (self.request_bucket.to_dict() if self.request_bucket else None),
+            "token_bucket": (self.token_bucket.to_dict() if self.token_bucket else None),
+        }
+
+
+class RateLimitedExecutor(Executor):
+    """Executor with integrated rate limiting via RateLimitedProcessor.
+
+    Manages processor lifecycle and forwards events for permission checking.
+
+    Example:
+        >>> executor = RateLimitedExecutor(processor_config={
+        ...     "queue_capacity": 50,
+        ...     "capacity_refresh_time": 60.0,
+        ...     "request_bucket": req_bucket,
+        ...     "token_bucket": tok_bucket,
+        ... })
+        >>> await executor.start()
+    """
+
+    processor_type = RateLimitedProcessor
+
+    def __init__(
+        self,
+        processor_config: dict[str, Any] | None = None,
+        strict_event_type: bool = False,
+        name: str | None = None,
+    ) -> None:
+        """Initialize rate-limited executor.
+
+        Args:
+            processor_config: Config dict for RateLimitedProcessor.create().
+            strict_event_type: If True, Flow enforces exact type matching.
+            name: Optional name for the executor Flow.
+        """
+        super().__init__(
+            processor_config=processor_config,
+            strict_event_type=strict_event_type,
+            name=name or "rate_limited_executor",
+        )
+
+    @override
+    async def start(self) -> None:
+        """Start executor and spawn replenishment task if not running."""
+        await super().start()
+
+        if (
+            self.processor
+            and isinstance(self.processor, RateLimitedProcessor)
+            and not self.processor._replenisher_task
+        ):
+            import asyncio
+
+            self.processor._replenisher_task = asyncio.create_task(
+                self.processor.start_replenishing()
+            )

kronos/services/utilities/rate_limiter.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Token bucket rate limiter for controlling request/token throughput."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+from kronos.utils.concurrency import Lock, current_time, sleep
+
+__all__ = ("RateLimitConfig", "TokenBucket")
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class RateLimitConfig:
+    """Immutable configuration for TokenBucket rate limiter.
+
+    Args:
+        capacity: Maximum tokens the bucket can hold.
+        refill_rate: Tokens added per second.
+        initial_tokens: Starting tokens (defaults to capacity).
+
+    Example:
+        >>> config = RateLimitConfig(capacity=100, refill_rate=10.0)
+        >>> bucket = TokenBucket(config)
+    """
+
+    capacity: int
+    refill_rate: float
+    initial_tokens: int | None = None
+
+    def __post_init__(self):
+        """Validate configuration parameters."""
+        if self.capacity <= 0:
+            raise ValueError("capacity must be > 0")
+        if self.refill_rate <= 0:
+            raise ValueError("refill_rate must be > 0")
+        if self.initial_tokens is None:
+            object.__setattr__(self, "initial_tokens", self.capacity)
+        elif self.initial_tokens < 0:
+            raise ValueError("initial_tokens must be >= 0")
+        elif self.initial_tokens > self.capacity:
+            raise ValueError(
+                f"initial_tokens ({self.initial_tokens}) cannot exceed capacity ({self.capacity})"
+            )
+
+
+class TokenBucket:
+    """Token bucket rate limiter with automatic refill.
+
+    Tokens are consumed on acquire() and refilled continuously based on
+    elapsed time. Thread-safe via async lock.
+
+    Example:
+        >>> config = RateLimitConfig(capacity=100, refill_rate=10.0)
+        >>> bucket = TokenBucket(config)
+        >>> if await bucket.try_acquire(5):
+        ...     # proceed with rate-limited operation
+        ...     pass
+
+    Attributes:
+        capacity: Maximum tokens the bucket can hold.
+        refill_rate: Tokens added per second.
+        tokens: Current available tokens (float for partial refills).
+    """
+
+    def __init__(self, config: RateLimitConfig):
+        """Initialize bucket from config."""
+        self.capacity = config.capacity
+        self.refill_rate = config.refill_rate
+        assert config.initial_tokens is not None
+        self.tokens = float(config.initial_tokens)
+        self.last_refill = current_time()
+        self._lock = Lock()
+
+    async def acquire(self, tokens: int = 1, *, timeout: float | None = None) -> bool:
+        """Acquire N tokens, waiting if necessary.
+
+        Args:
+            tokens: Number of tokens to acquire
+            timeout: Max wait time in seconds (None = wait forever)
+
+        Returns:
+            True if acquired, False if timeout
+
+        Raises:
+            ValueError: If tokens <= 0 or tokens > capacity
+        """
+        if tokens <= 0:
+            raise ValueError("tokens must be > 0")
+        if tokens > self.capacity:
+            raise ValueError(
+                f"Cannot acquire {tokens} tokens: exceeds bucket capacity {self.capacity}"
+            )
+
+        start_time = current_time()
+
+        while True:
+            async with self._lock:
+                self._refill()
+
+                if self.tokens >= tokens:
+                    self.tokens -= tokens
+                    logger.debug(f"Acquired {tokens} tokens, {self.tokens:.2f} remaining")
+                    return True
+
+                deficit = tokens - self.tokens
+                wait_time = deficit / self.refill_rate
+
+            # Check timeout
+            if timeout is not None:
+                elapsed = current_time() - start_time
+                if elapsed + wait_time > timeout:
+                    logger.warning(f"Rate limit timeout after {elapsed:.2f}s")
+                    return False
+                wait_time = min(wait_time, timeout - elapsed)
+
+            logger.debug(f"Waiting {wait_time:.2f}s for {deficit:.2f} tokens")
+            await sleep(wait_time)
+
+    def _refill(self) -> None:
+        """Refill tokens based on elapsed time (call under lock)."""
+        now = current_time()
+        elapsed = now - self.last_refill
+        new_tokens = elapsed * self.refill_rate
+
+        self.tokens = min(self.capacity, self.tokens + new_tokens)
+        self.last_refill = now
+
+    async def try_acquire(self, tokens: int = 1) -> bool:
+        """Try to acquire tokens without waiting.
+
+        Returns:
+            True if acquired immediately, False if insufficient tokens
+
+        Raises:
+            ValueError: If tokens <= 0
+        """
+        if tokens <= 0:
+            raise ValueError("tokens must be > 0")
+        async with self._lock:
+            self._refill()
+
+            if self.tokens >= tokens:
+                self.tokens -= tokens
+                return True
+            return False
+
+    async def reset(self) -> None:
+        """Reset bucket to full capacity (thread-safe).
+
+        Used by RateLimitedProcessor for interval-based replenishment.
+        """
+        async with self._lock:
+            self.tokens = float(self.capacity)
+            self.last_refill = current_time()
+
+    async def release(self, tokens: int = 1) -> None:
+        """Release tokens back to bucket (thread-safe).
+
+        Used for rollback when dual-bucket acquire fails partway through.
+
+        Args:
+            tokens: Number of tokens to release back.
+
+        Raises:
+            ValueError: If tokens <= 0.
+        """
+        if tokens <= 0:
+            raise ValueError("tokens must be > 0")
+        async with self._lock:
+            self.tokens = min(self.capacity, self.tokens + tokens)
+
+    def to_dict(self) -> dict[str, float]:
+        """Serialize config to dict (excludes runtime state)."""
+        return {"capacity": self.capacity, "refill_rate": self.refill_rate}