keymesh 0.1.2a0__py3-none-any.whl

keymesh/__init__.py

@@ -0,0 +1,45 @@
+"""
+KeyMesh — Credential Orchestration Runtime for AI API Systems.
+
+A lightweight, concurrency-safe API key pool manager and scheduler.
+KeyMesh is NOT a proxy, SDK wrapper, or HTTP gateway.
+It is purely a credential allocator and rate-limit-aware scheduler.
+
+Usage:
+    from keymesh import KeyPool, SchedulerStrategy
+
+    pool = KeyPool(
+        keys=["sk-key1", "sk-key2", "sk-key3"],
+        strategy=SchedulerStrategy.LEAST_BUSY,
+    )
+
+    key = await pool.acquire()
+    # ... use key in your own SDK/HTTP client ...
+    await pool.release(key)
+"""
+
+from keymesh.pool.pool import KeyPool
+from keymesh.pool.sync_pool import SyncKeyPool
+from keymesh.scheduler.base import SchedulerStrategy
+from keymesh.state.key_state import KeyState
+from keymesh.state.sync_key_state import SyncKeyState
+from keymesh.exceptions import (
+    KeyMeshError,
+    NoAvailableKeyError,
+    KeyCooldownError,
+    KeyExhaustedError,
+)
+
+__all__ = [
+    "KeyPool",
+    "SyncKeyPool",
+    "SchedulerStrategy",
+    "KeyState",
+    "SyncKeyState",
+    "KeyMeshError",
+    "NoAvailableKeyError",
+    "KeyCooldownError",
+    "KeyExhaustedError",
+]
+
+__version__ = "0.1.0"

keymesh/concurrency/__init__.py

@@ -0,0 +1,4 @@
+"""keymesh.concurrency package."""
+from keymesh.concurrency.semaphores import SemaphoreGroup
+
+__all__ = ["SemaphoreGroup"]

keymesh/concurrency/semaphores.py

@@ -0,0 +1,47 @@
+"""
+Concurrency utilities for KeyMesh.
+
+Provides a SemaphoreGroup for per-key concurrency limits — optional,
+for users who want to cap in-flight requests per key.
+"""
+
+import asyncio
+
+
+class SemaphoreGroup:
+    """
+    Per-key asyncio semaphore pool.
+
+    Use this when you want to cap the number of concurrent in-flight
+    requests per API key (e.g., a provider that allows max 5 concurrent calls).
+
+    Semaphores are created lazily on first use inside a running event loop,
+    avoiding "Future attached to a different loop" errors when this object
+    is instantiated at module level or before ``asyncio.run()`` is called.
+
+    Example
+    -------
+    ```python
+    sem_group = SemaphoreGroup(max_concurrent=5)
+
+    key = await pool.acquire()
+    async with sem_group.acquire(key):
+        response = await client.call(api_key=key)
+    await pool.release(key)
+    ```
+    """
+
+    def __init__(self, max_concurrent: int = 10) -> None:
+        self._max = max_concurrent
+        self._sems: dict[str, asyncio.Semaphore] = {}
+
+    def acquire(self, key: str) -> asyncio.Semaphore:
+        """Return (or lazily create) the semaphore for `key`.
+
+        Must be called from within a running async context so that the
+        semaphore is bound to the correct event loop.
+        """
+        if key not in self._sems:
+            # Created inside a coroutine: the running loop is guaranteed.
+            self._sems[key] = asyncio.Semaphore(self._max)
+        return self._sems[key]

keymesh/cooldown/__init__.py

@@ -0,0 +1,4 @@
+"""keymesh.cooldown package."""
+from keymesh.cooldown.manager import CooldownManager
+
+__all__ = ["CooldownManager"]

keymesh/cooldown/manager.py

@@ -0,0 +1,52 @@
+"""
+Cooldown manager — tracks and enforces per-key cooldown windows.
+
+KeyMesh never sleeps waiting for cooldowns. Instead, unavailable keys
+are filtered out at acquire-time and re-admitted when time has elapsed.
+"""
+
+import time
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from keymesh.state.key_state import KeyState
+
+
+class CooldownManager:
+    """
+    Stateless helper for cooldown evaluation.
+
+    All actual cooldown state lives inside KeyState.
+    This class provides higher-level helpers used by the pool.
+    """
+
+    @staticmethod
+    def apply(key_state: "KeyState", duration: float = 60.0) -> None:
+        """
+        Immediately apply a cooldown window to a key (synchronous, for internal use).
+
+        .. deprecated::
+            This method mutates ``cooldown_until`` directly without acquiring
+            the key's internal ``asyncio.Lock``, making it unsafe under
+            concurrency. Use ``await key_state.apply_cooldown(duration)``
+            in async contexts instead.
+        """
+        import warnings
+
+        warnings.warn(
+            "CooldownManager.apply() is not concurrency-safe. "
+            "Use `await key_state.apply_cooldown(duration)` in async code.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        key_state.cooldown_until = time.monotonic() + duration
+
+    @staticmethod
+    def is_expired(key_state: "KeyState") -> bool:
+        """Return True if cooldown has elapsed (key may be re-admitted)."""
+        return time.monotonic() >= key_state.cooldown_until
+
+    @staticmethod
+    def remaining(key_state: "KeyState") -> float:
+        """Seconds until cooldown expires. Returns 0 if not cooling down."""
+        return max(0.0, key_state.cooldown_until - time.monotonic())

keymesh/exceptions.py

@@ -0,0 +1,30 @@
+"""
+KeyMesh custom exceptions.
+"""
+
+
+class KeyMeshError(Exception):
+    """Base exception for all KeyMesh errors."""
+
+
+class NoAvailableKeyError(KeyMeshError):
+    """Raised when no API key is currently available (all cooling down or failed)."""
+
+
+class KeyCooldownError(KeyMeshError):
+    """Raised when a specific key is in cooldown and cannot be used."""
+
+    def __init__(self, key: str, cooldown_remaining: float) -> None:
+        self.key = key
+        self.cooldown_remaining = cooldown_remaining
+        super().__init__(
+            f"Key ...{key[-6:]} is in cooldown for {cooldown_remaining:.1f}s more."
+        )
+
+
+class KeyExhaustedError(KeyMeshError):
+    """Raised when a key has exceeded its failure threshold and is removed from the pool."""
+
+    def __init__(self, key: str) -> None:
+        self.key = key
+        super().__init__(f"Key ...{key[-6:]} has been exhausted and removed from pool.")

keymesh/metrics/__init__.py

@@ -0,0 +1,4 @@
+"""keymesh.metrics package."""
+from keymesh.metrics.pool_metrics import PoolMetrics
+
+__all__ = ["PoolMetrics"]

keymesh/metrics/pool_metrics.py

@@ -0,0 +1,52 @@
+"""
+Runtime metrics aggregation for the pool.
+"""
+
+import time
+from dataclasses import dataclass, field
+
+
+@dataclass
+class PoolMetrics:
+    """Aggregated runtime metrics for a KeyPool instance."""
+
+    total_acquires: int = 0
+    total_releases: int = 0
+    total_failures: int = 0
+    total_cooldowns: int = 0
+    total_retries: int = 0
+    no_key_available_count: int = 0
+    _started_at: float = field(default_factory=time.monotonic, init=False)
+
+    def record_acquire(self) -> None:
+        self.total_acquires += 1
+
+    def record_release(self) -> None:
+        self.total_releases += 1
+
+    def record_failure(self) -> None:
+        self.total_failures += 1
+
+    def record_cooldown(self) -> None:
+        self.total_cooldowns += 1
+
+    def record_retry(self) -> None:
+        self.total_retries += 1
+
+    def record_no_key(self) -> None:
+        self.no_key_available_count += 1
+
+    @property
+    def uptime_seconds(self) -> float:
+        return time.monotonic() - self._started_at
+
+    def snapshot(self) -> dict[str, object]:
+        return {
+            "uptime_seconds": round(self.uptime_seconds, 2),
+            "total_acquires": self.total_acquires,
+            "total_releases": self.total_releases,
+            "total_failures": self.total_failures,
+            "total_cooldowns": self.total_cooldowns,
+            "total_retries": self.total_retries,
+            "no_key_available_count": self.no_key_available_count,
+        }

keymesh/pool/__init__.py

@@ -0,0 +1,5 @@
+"""keymesh.pool package."""
+from keymesh.pool.pool import KeyPool
+from keymesh.pool.sync_pool import SyncKeyPool
+
+__all__ = ["KeyPool", "SyncKeyPool"]

keymesh/pool/pool.py

@@ -0,0 +1,251 @@
+"""
+KeyPool — the central public interface of KeyMesh.
+
+Responsibilities:
+  - Hold and manage a pool of KeyState objects
+  - Delegate key selection to the configured scheduler
+  - Enforce cooldowns, failure thresholds, and retry logic
+  - Surface a minimal public API: acquire / release / mark_failed / mark_rate_limited
+
+KeyPool is framework-agnostic. It never sends HTTP requests or wraps any SDK.
+"""
+
+import asyncio
+import logging
+from typing import Any, Sequence
+
+from keymesh.exceptions import KeyExhaustedError, NoAvailableKeyError
+from keymesh.metrics.pool_metrics import PoolMetrics
+from keymesh.scheduler import BaseScheduler, SchedulerStrategy, build_scheduler
+from keymesh.state.key_state import KeyState
+from keymesh.storage.base import BaseStorage
+from keymesh.storage.memory import MemoryStorage
+
+logger = logging.getLogger(__name__)
+
+
+class KeyPool:
+    """
+    Manages a collection of API keys with scheduling, cooldown, and retry support.
+
+    Example
+    -------
+    ```python
+    pool = KeyPool(keys=["sk-a", "sk-b", "sk-c"])
+
+    key = await pool.acquire()
+    try:
+        response = await my_client.call(api_key=key)
+        await pool.release(key, latency=response.elapsed)
+    except RateLimitError:
+        await pool.mark_rate_limited(key, cooldown=60.0)
+    except Exception:
+        await pool.mark_failed(key)
+    ```
+    """
+
+    def __init__(
+        self,
+        keys: Sequence[str],
+        *,
+        strategy: SchedulerStrategy = SchedulerStrategy.LEAST_BUSY,
+        scheduler: BaseScheduler | None = None,
+        storage: BaseStorage | None = None,
+        default_cooldown: float = 60.0,
+        max_failures: int = 10,
+        acquire_timeout: float = 30.0,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        keys:
+            List of raw API key strings. Must be non-empty.
+        strategy:
+            Scheduling strategy to use. Ignored if `scheduler` is provided.
+        scheduler:
+            Custom scheduler instance. Overrides `strategy`.
+        storage:
+            Persistence backend. Defaults to in-memory (no persistence).
+        default_cooldown:
+            Default cooldown duration in seconds when mark_rate_limited is called.
+        max_failures:
+            Consecutive failures before a key is considered exhausted.
+        acquire_timeout:
+            Max seconds to wait for an available key before raising NoAvailableKeyError.
+        """
+        if not keys:
+            raise ValueError("KeyPool requires at least one API key.")
+
+        self._states: dict[str, KeyState] = {
+            k: KeyState(key=k, max_failures=max_failures) for k in keys
+        }
+        self._scheduler: BaseScheduler = scheduler or build_scheduler(strategy)
+        self._storage: BaseStorage = storage or MemoryStorage()
+        self._default_cooldown = default_cooldown
+        self._max_failures = max_failures
+        self._acquire_timeout = acquire_timeout
+        self._metrics = PoolMetrics()
+        self._pool_lock = asyncio.Lock()
+
+        logger.info(
+            "KeyPool initialised: %d keys, strategy=%s",
+            len(keys),
+            self._scheduler.strategy,
+        )
+
+    # ── Public API ────────────────────────────────────────────────────────────────
+
+    async def acquire(self) -> str:
+        """
+        Acquire an available API key.
+
+        Blocks (up to `acquire_timeout`) until a key becomes available.
+
+        Returns
+        -------
+        str
+            The raw API key string to inject into your HTTP client / SDK.
+
+        Raises
+        ------
+        NoAvailableKeyError
+            If no key becomes available within `acquire_timeout` seconds.
+        """
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + self._acquire_timeout
+
+        while True:
+            key_state = await self._try_acquire()
+            if key_state is not None:
+                self._metrics.record_acquire()
+                logger.debug("Acquired key ...%s", key_state.key[-6:])
+                return key_state.key
+
+            self._metrics.record_no_key()
+            remaining = deadline - loop.time()
+            if remaining <= 0:
+                raise NoAvailableKeyError(
+                    "No API keys are currently available. "
+                    "All keys are either cooling down or exhausted."
+                )
+
+            # Brief sleep before retrying to avoid busy-loop
+            await asyncio.sleep(min(0.5, remaining))
+
+    async def release(self, key: str, *, latency: float = 0.0) -> None:
+        """
+        Release a key after a successful API call.
+
+        Parameters
+        ----------
+        key:
+            The key string returned by acquire().
+        latency:
+            Optional: elapsed seconds for the call (used for weighted scheduling).
+        """
+        state = self._resolve(key)
+        await state.record_success(latency)
+        self._metrics.record_release()
+        await self._storage.save(key, state.snapshot())
+        logger.debug("Released key ...%s (latency=%.3fs)", key[-6:], latency)
+
+    async def mark_failed(self, key: str) -> None:
+        """
+        Mark a key as failed (non-rate-limit error).
+
+        After `max_failures` calls, the key is marked exhausted and excluded
+        from future scheduling.
+
+        Raises
+        ------
+        KeyExhaustedError
+            If the key has now reached the failure threshold.
+        """
+        state = self._resolve(key)
+        await state.record_failure()
+        self._metrics.record_failure()
+        await self._storage.save(key, state.snapshot())
+        logger.warning("Key ...%s marked failed (total=%d)", key[-6:], state.failure_count)
+
+        if state.is_exhausted:
+            raise KeyExhaustedError(key)
+
+    async def mark_rate_limited(self, key: str, *, cooldown: float | None = None) -> None:
+        """
+        Mark a key as rate-limited (HTTP 429).
+
+        The key will be excluded from scheduling until the cooldown elapses.
+        Resets the key's consecutive failure counter — a rate-limited key is
+        still valid, just throttled.
+
+        Parameters
+        ----------
+        cooldown:
+            Cooldown duration in seconds. Defaults to pool's `default_cooldown`.
+        """
+        duration = cooldown if cooldown is not None else self._default_cooldown
+        state = self._resolve(key)
+        await state.apply_cooldown(duration)
+        self._metrics.record_cooldown()
+        await self._storage.save(key, state.snapshot())
+        logger.warning(
+            "Key ...%s rate-limited, cooling down for %.1fs", key[-6:], duration
+        )
+
+    async def add_key(self, key: str) -> None:
+        """
+        Add or reset a key in the pool at runtime.
+
+        Use this to:
+        - Re-admit an exhausted key after rotating/refreshing it.
+        - Inject a brand-new key into a running pool without restarting.
+
+        Parameters
+        ----------
+        key:
+            The raw API key string to add or reset.
+        """
+        async with self._pool_lock:
+            self._states[key] = KeyState(key=key, max_failures=self._max_failures)
+        logger.info("Key ...%s added/reset in pool", key[-6:])
+
+    # ── Diagnostics ───────────────────────────────────────────────────────────────
+
+    def status(self) -> dict[str, Any]:
+        """Return a full status snapshot of the pool and all key states."""
+        return {
+            "pool_metrics": self._metrics.snapshot(),
+            "scheduler": self._scheduler.strategy,
+            "keys": [ks.snapshot() for ks in self._states.values()],
+        }
+
+    def available_count(self) -> int:
+        """Number of keys currently available for scheduling."""
+        return sum(1 for ks in self._states.values() if ks.is_available)
+
+    # ── Internal ──────────────────────────────────────────────────────────────────
+
+    async def _try_acquire(self) -> KeyState | None:
+        """
+        Ask the scheduler to pick one available key.
+
+        Returns None if no key is available right now.
+        """
+        async with self._pool_lock:
+            candidates = [ks for ks in self._states.values() if ks.is_available]
+            key_state = self._scheduler.select(candidates)
+            if key_state is not None:
+                await key_state.increment_active()
+            return key_state
+
+    def _resolve(self, key: str) -> KeyState:
+        """Look up a KeyState by raw key string. Raises KeyError if unknown."""
+        try:
+            return self._states[key]
+        except KeyError as exc:
+            raise KeyError(f"Unknown key '...{key[-6:]}' — was it added to this pool?") from exc
+
+    async def close(self) -> None:
+        """Release storage resources. Call when pool is no longer needed."""
+        await self._storage.close()
+        logger.info("KeyPool closed.")