eggpool 0.1.0__py3-none-any.whl

eggpool/accounts/registry.py

@@ -0,0 +1,181 @@
+"""Account registry: loads accounts from config, manages runtime state."""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING, Any
+
+from eggpool.accounts.state import AccountRuntimeState
+from eggpool.errors import ConfigError
+
+if TYPE_CHECKING:
+    from eggpool.models.config import AppConfig
+
+logger = logging.getLogger(__name__)
+
+
+def account_config_rows(config: AppConfig) -> list[dict[str, Any]]:
+    """Serialize configured accounts into rows for persistence.
+
+    Returns a list of dicts with the exact fields consumed by
+    :meth:`eggpool.db.repositories.AccountRepository.sync_from_config`.
+    Keeping the shape in one place prevents the app lifespan and the
+    ``models refresh`` CLI command from drifting out of sync.
+    """
+    rows: list[dict[str, Any]] = []
+    for provider_id, provider in config.providers.items():
+        for acct in provider.accounts:
+            rows.append(
+                {
+                    "name": acct.name,
+                    "api_key_env": acct.api_key_env,
+                    "enabled": acct.enabled,
+                    "weight": acct.weight,
+                    "provider_id": provider_id,
+                }
+            )
+    return rows
+
+
+class AccountRegistry:
+    """Manages account configurations and their runtime states."""
+
+    def __init__(self, config: AppConfig) -> None:
+        self._config = config
+        self._states: dict[str, AccountRuntimeState] = {}
+        self._api_keys: dict[str, str] = {}
+        self._account_providers: dict[str, str] = {}
+        self._initialize()
+
+    def _initialize(self) -> None:
+        """Load accounts from config and resolve API keys."""
+        for provider_id, provider_cfg in self._config.providers.items():
+            for acct_config in provider_cfg.accounts:
+                api_key = acct_config.api_key or os.environ.get(
+                    acct_config.api_key_env, ""
+                )
+                if (
+                    acct_config.enabled
+                    and provider_cfg.auth.mode != "none"
+                    and not api_key
+                ):
+                    raise ConfigError(
+                        f"Account {acct_config.name!r} is enabled but "
+                        f"no API key available"
+                    )
+
+                state = AccountRuntimeState(
+                    name=acct_config.name,
+                    enabled=acct_config.enabled,
+                    weight=acct_config.weight,
+                    routing_priority=provider_cfg.routing_priority,
+                )
+                self._states[acct_config.name] = state
+                self._api_keys[acct_config.name] = api_key
+                self._account_providers[acct_config.name] = provider_id
+
+                if acct_config.enabled:
+                    logger.info(
+                        "Loaded account %r (weight=%.2f, priority=%d, provider=%r)",
+                        acct_config.name,
+                        acct_config.weight,
+                        provider_cfg.routing_priority,
+                        provider_id,
+                    )
+
+    def reload(self, config: AppConfig) -> None:
+        """Reload account configurations from a new config."""
+        self._config = config
+        self._states.clear()
+        self._api_keys.clear()
+        self._account_providers.clear()
+        self._initialize()
+
+    def get_state(self, name: str) -> AccountRuntimeState | None:
+        """Get runtime state for an account by name."""
+        return self._states.get(name)
+
+    def get_api_key(self, name: str) -> str | None:
+        """Get the resolved API key for an account."""
+        return self._api_keys.get(name)
+
+    def has_usable_credentials(self, name: str) -> bool:
+        """Return whether an account can satisfy its provider auth contract."""
+        provider_id = self._account_providers.get(name)
+        if provider_id is None:
+            return False
+        provider = self._config.providers.get(provider_id)
+        if provider is None:
+            return False
+        return provider.auth.mode == "none" or bool(self._api_keys.get(name))
+
+    def get_all_states(self) -> list[AccountRuntimeState]:
+        """Get all account runtime states."""
+        return list(self._states.values())
+
+    def get_enabled_states(self) -> list[AccountRuntimeState]:
+        """Get runtime states for enabled accounts."""
+        return [s for s in self._states.values() if s.enabled]
+
+    def get_eligible_states(self) -> list[AccountRuntimeState]:
+        """Get runtime states for eligible accounts."""
+        return [s for s in self._states.values() if s.is_eligible()]
+
+    def get_account_config(self, name: str):
+        """Get the config for an account by name."""
+        for acct in self._config.all_accounts():
+            if acct.name == name:
+                return acct
+        return None
+
+    def get_account_offsets(self, name: str) -> dict[str, int]:
+        """Get quota offsets for an account."""
+        acct = self.get_account_config(name)
+        if acct is None:
+            return {}
+        return {
+            "five_hour": acct.five_hour_offset_microdollars,
+            "weekly": acct.weekly_offset_microdollars,
+            "monthly": acct.monthly_offset_microdollars,
+        }
+
+    def get_provider_for_account(self, account_name: str) -> str | None:
+        """Get the provider ID for an account."""
+        return self._account_providers.get(account_name)
+
+    def get_provider_protocols(self, provider_id: str) -> set[str]:
+        """Get protocols configured for a provider."""
+        provider = self._config.providers.get(provider_id)
+        if provider is None:
+            return set()
+        return set(provider.protocols)
+
+    def account_supports_protocol(self, account_name: str, protocol: str) -> bool:
+        """Return whether an account's configured provider supports a protocol."""
+        provider_id = self.get_provider_for_account(account_name)
+        if provider_id is None:
+            return False
+        return protocol in self.get_provider_protocols(provider_id)
+
+    def get_accounts_for_provider(self, provider_id: str) -> list[AccountRuntimeState]:
+        """Get all account states belonging to a provider."""
+        return [
+            state
+            for name, state in self._states.items()
+            if self._account_providers.get(name) == provider_id
+        ]
+
+    def get_enabled_accounts_for_provider(
+        self, provider_id: str
+    ) -> list[AccountRuntimeState]:
+        """Get enabled account states belonging to a provider."""
+        return [
+            state
+            for name, state in self._states.items()
+            if self._account_providers.get(name) == provider_id and state.enabled
+        ]
+
+    def get_provider_ids(self) -> list[str]:
+        """Get all unique provider IDs."""
+        return sorted(set(self._account_providers.values()))

eggpool/accounts/state.py

@@ -0,0 +1,170 @@
+"""Account runtime state and health management."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+
+# Default cooldown durations used when the caller does not supply an
+# explicit value. These mirror the configured defaults in
+# ``RoutingConfig`` and exist only as a safety net for tests or
+# command-line tools that instantiate ``AccountRuntimeState``
+# directly. Production code paths should pass the configured cooldown
+# explicitly so the runtime state stays in lock-step with the
+# authoritative ``HealthManager``.
+DEFAULT_QUOTA_EXHAUSTED_COOLDOWN_SECONDS = 300.0
+DEFAULT_BACKOFF_BASE_SECONDS = 30.0
+DEFAULT_BACKOFF_MAX_SECONDS = 3600.0  # 1 hour max backoff for rate limits
+
+
+def _failure_backoff(consecutive_failures: int) -> float:
+    """Return capped exponential backoff without constructing huge integers."""
+    if consecutive_failures <= 1:
+        return DEFAULT_BACKOFF_BASE_SECONDS
+    max_doublings = int(
+        DEFAULT_BACKOFF_MAX_SECONDS / DEFAULT_BACKOFF_BASE_SECONDS
+    ).bit_length()
+    doublings = min(consecutive_failures - 1, max_doublings)
+    return min(
+        DEFAULT_BACKOFF_BASE_SECONDS * (2**doublings),
+        DEFAULT_BACKOFF_MAX_SECONDS,
+    )
+
+
+@dataclass
+class AccountRuntimeState:
+    """Mutable runtime state for an account."""
+
+    name: str
+    enabled: bool = True
+    weight: float = 1.0
+    routing_priority: int = 0
+
+    health_state: str = "healthy"
+    cooldown_until: float = 0.0
+    consecutive_failures: int = 0
+    last_success_at: float = 0.0
+    last_failure_at: float = 0.0
+    last_failure_category: str = ""
+
+    active_request_count: int = 0
+    reserved_microdollars: int = 0
+
+    # Account-specific model availability: model_id -> available
+    model_availability: dict[str, bool] = field(  # pyright: ignore[reportUnknownVariableType]
+        default_factory=dict
+    )
+
+    def refresh_transient_state(self, now: float | None = None) -> None:
+        """Clear transient cooldown status when it expires.
+
+        Auto-recoverable states ("rate_limited", "cooldown") are
+        cleared either when the configured cooldown has elapsed or
+        when no cooldown is set. "quota_exhausted" recovers
+        automatically when the cooldown expires.
+        """
+        if now is None:
+            now = time.time()
+        if self.health_state == "quota_exhausted":
+            if self.cooldown_until > 0 and now >= self.cooldown_until:
+                self.health_state = "healthy"
+                self.cooldown_until = 0.0
+                self.consecutive_failures = 0
+            return
+        if self.health_state in ("rate_limited", "cooldown") and (
+            self.cooldown_until == 0.0 or now >= self.cooldown_until
+        ):
+            self.health_state = "healthy"
+            self.cooldown_until = 0.0
+            self.consecutive_failures = 0
+
+    def is_eligible(self) -> bool:
+        """Check if account is eligible for routing."""
+        if not self.enabled:
+            return False
+        self.refresh_transient_state()
+        if self.health_state in (
+            "authentication_failed",
+            "quota_exhausted",
+            "cooldown",
+            "rate_limited",
+        ):
+            return False
+        return self.cooldown_until <= time.time()
+
+    def record_success(self) -> None:
+        """Record a successful request."""
+        self.consecutive_failures = 0
+        self.last_success_at = time.time()
+        self.last_failure_category = ""
+        if self.health_state in ("cooldown", "rate_limited", "quota_exhausted"):
+            self.health_state = "healthy"
+            self.cooldown_until = 0.0
+
+    def record_failure(
+        self,
+        error_class: str,
+        *,
+        cooldown_seconds: float | None = None,
+        rate_limit_retry_after: float | None = None,
+    ) -> None:
+        """Record a failed request and update health state.
+
+        ``cooldown_seconds`` is the configured quota-exhausted cooldown
+        duration; the same value used by the authoritative
+        ``HealthManager`` must be passed here so the two cooldown
+        representations cannot diverge. ``rate_limit_retry_after`` is
+        the parsed ``Retry-After`` value for 429 responses; when
+        supplied, it takes precedence over the exponential backoff
+        schedule. Authentication failures remain terminal until
+        explicitly reset.
+        """
+        self.consecutive_failures += 1
+        self.last_failure_at = time.time()
+
+        if self.last_failure_category and error_class != self.last_failure_category:
+            self.consecutive_failures = 1
+        self.last_failure_category = error_class
+
+        if error_class in ("authentication_failed", "authentication"):
+            self.health_state = "authentication_failed"
+        elif error_class == "quota_exhausted":
+            self.health_state = "quota_exhausted"
+            duration = (
+                cooldown_seconds
+                if cooldown_seconds is not None
+                else DEFAULT_QUOTA_EXHAUSTED_COOLDOWN_SECONDS
+            )
+            self.cooldown_until = time.time() + duration
+        elif error_class == "rate_limited":
+            # Mirror HealthManager.record_rate_limit so both state
+            # machines expose the same label for the same event.
+            self.health_state = "rate_limited"
+            if rate_limit_retry_after is not None:
+                self.cooldown_until = time.time() + max(0.0, rate_limit_retry_after)
+            else:
+                self.cooldown_until = time.time() + _failure_backoff(
+                    self.consecutive_failures
+                )
+        elif error_class in (
+            "connect_timeout",
+            "read_timeout",
+            "connection_failure",
+            "connection_error",
+        ):
+            self.health_state = "cooldown"
+            self.cooldown_until = time.time() + _failure_backoff(
+                self.consecutive_failures
+            )
+        # upstream_server_error, protocol_error, unknown, etc. - no cooldown
+
+    def reset_health(self) -> None:
+        """Reset health state to healthy."""
+        self.health_state = "healthy"
+        self.cooldown_until = 0.0
+        self.consecutive_failures = 0
+        self.last_failure_category = ""
+        # Per-model disable map must be cleared alongside the
+        # account-level state so reset_health mirrors the same
+        # reset semantics as HealthManager.enable_account.
+        self.model_availability.clear()

eggpool/api/__init__.py

@@ -0,0 +1 @@
+"""API endpoint handlers."""

eggpool/api/chat_completions.py

@@ -0,0 +1,31 @@
+"""OpenAI-compatible ``/v1/chat/completions`` endpoint."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from fastapi import Request  # noqa: TCH002 — FastAPI needs runtime access
+
+from eggpool.api.errors import openai_error_response
+from eggpool.api.proxy_request import (
+    ProxyEndpointConfig,
+    handle_proxy_request,
+)
+
+if TYPE_CHECKING:
+    from fastapi.responses import Response
+
+_ENDPOINT = ProxyEndpointConfig(
+    protocol="openai",
+    request_label="chat completion",
+    error_response=openai_error_response,
+    not_found_error_type="invalid_request_error",
+    service_error_type="server_error",
+)
+
+
+async def handle_chat_completions(
+    request: Request,
+) -> Response:
+    """Handle POST /v1/chat/completions."""
+    return await handle_proxy_request(request, _ENDPOINT)

eggpool/api/errors.py

@@ -0,0 +1,51 @@
+"""Protocol-specific error response renderers.
+
+Provides OpenAI-style and Anthropic-style error response formats
+so that upstream clients receive protocol-compatible error payloads.
+"""
+
+from __future__ import annotations
+
+from fastapi.responses import JSONResponse
+
+
+def openai_error_response(
+    status_code: int,
+    message: str,
+    error_type: str = "invalid_request_error",
+) -> JSONResponse:
+    """Return an OpenAI-compatible error response."""
+    return JSONResponse(
+        status_code=status_code,
+        content={
+            "error": {
+                "message": message,
+                "type": error_type,
+                "code": str(status_code),
+            }
+        },
+    )
+
+
+def anthropic_error_response(
+    status_code: int,
+    message: str,
+    error_type: str = "invalid_request_error",
+) -> JSONResponse:
+    """Return an Anthropic-compatible error response."""
+    return JSONResponse(
+        status_code=status_code,
+        content={
+            "type": "error",
+            "error": {
+                "type": error_type,
+                "message": message,
+            },
+        },
+    )
+
+
+__all__ = [
+    "anthropic_error_response",
+    "openai_error_response",
+]

eggpool/api/messages.py

@@ -0,0 +1,31 @@
+"""Anthropic-compatible ``/v1/messages`` endpoint."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from fastapi import Request  # noqa: TCH002 — FastAPI needs runtime access
+
+from eggpool.api.errors import anthropic_error_response
+from eggpool.api.proxy_request import (
+    ProxyEndpointConfig,
+    handle_proxy_request,
+)
+
+if TYPE_CHECKING:
+    from fastapi.responses import Response
+
+_ENDPOINT = ProxyEndpointConfig(
+    protocol="anthropic",
+    request_label="messages request",
+    error_response=anthropic_error_response,
+    not_found_error_type="not_found_error",
+    service_error_type="api_error",
+)
+
+
+async def handle_messages(
+    request: Request,
+) -> Response:
+    """Handle POST /v1/messages."""
+    return await handle_proxy_request(request, _ENDPOINT)

eggpool/api/models.py

@@ -0,0 +1,76 @@
+"""Serialization helpers for /v1/models responses."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
+def serialize_openai_model(
+    model: Mapping[str, Any],
+    *,
+    routing_priority: int | None = None,
+    routing_priority_max: int | None = None,
+    providers: list[str] | None = None,
+) -> dict[str, Any]:
+    """Serialize a catalog model entry to OpenAI-compatible model dict.
+
+    Includes the namespaced ``eggpool`` extension with base model ID,
+    provider ID, routing priority (when supplied), and effective limits
+    when available.
+
+    For collapsed entries (no per-provider ``provider_id``), pass the
+    ``routing_priority_max`` (highest priority across contributing
+    providers) and ``providers`` list so clients can see the routing
+    topology.
+    """
+    result: dict[str, Any] = {
+        "id": model["model_id"],
+        "object": "model",
+        "created": int(model.get("first_seen_at", 0)),
+        "owned_by": model.get("provider_id", "opencode"),
+        "name": model.get("display_name") or model["model_id"],
+    }
+
+    # Add namespaced EggPool metadata
+    eggpool_meta: dict[str, Any] = {}
+    base_model_id = model.get("base_model_id")
+    provider_id = model.get("provider_id")
+    if base_model_id is not None:
+        eggpool_meta["base_model_id"] = base_model_id
+    if provider_id is not None:
+        eggpool_meta["provider_id"] = provider_id
+    if routing_priority is not None:
+        eggpool_meta["routing_priority"] = routing_priority
+
+    # Collapsed-entry metadata: contributing providers and the highest
+    # routing priority across them. Both are omitted when the entry is
+    # already provider-scoped (the singular `routing_priority` above
+    # covers that case).
+    if provider_id is None:
+        if providers is not None:
+            eggpool_meta["providers"] = list(providers)
+        if routing_priority_max is not None:
+            eggpool_meta["routing_priority_max"] = routing_priority_max
+
+    effective = model.get("effective_limits", {})
+    if effective:
+        limits: dict[str, Any] = {}
+        ctx = effective.get("context_tokens")
+        inp = effective.get("input_tokens")
+        out = effective.get("output_tokens")
+        if ctx is not None:
+            limits["context"] = ctx
+        if inp is not None:
+            limits["input"] = inp
+        if out is not None:
+            limits["output"] = out
+        if limits:
+            eggpool_meta["limits"] = limits
+
+    if eggpool_meta:
+        result["eggpool"] = eggpool_meta
+
+    return result