llm-pycascade 0.1.0__py3-none-any.whl

llm_pycascade/__init__.py

@@ -0,0 +1,47 @@
+"""llm-pycascade — Resilient cascading LLM inference
+with failover and circuit breaking."""
+
+from llm_pycascade.cascade import run_cascade
+from llm_pycascade.config import (
+    AppConfig,
+    CascadeConfig,
+    CascadeEntry,
+    DatabaseConfig,
+    FailureConfig,
+    ProviderConfig,
+    ProviderType,
+    load_config,
+)
+from llm_pycascade.db import init_db
+from llm_pycascade.error import CascadeError, ProviderError
+from llm_pycascade.models import (
+    ContentBlock,
+    Conversation,
+    LlmResponse,
+    Message,
+    MessageRole,
+    ToolDefinition,
+)
+
+__all__ = [
+    "run_cascade",
+    "AppConfig",
+    "CascadeConfig",
+    "CascadeEntry",
+    "DatabaseConfig",
+    "FailureConfig",
+    "ProviderConfig",
+    "ProviderType",
+    "load_config",
+    "init_db",
+    "CascadeError",
+    "ProviderError",
+    "ContentBlock",
+    "Conversation",
+    "LlmResponse",
+    "Message",
+    "MessageRole",
+    "ToolDefinition",
+]
+
+__version__ = "0.1.0"

llm_pycascade/cascade.py

@@ -0,0 +1,287 @@
+"""Cascade engine — resilient multi-provider LLM inference with failover.
+
+The cascade iterates through an ordered list of provider/model entries.
+For each entry, it checks whether the combination is on cooldown, attempts
+the request, and on failure computes an exponential-backoff cooldown
+before moving to the next entry.
+
+If every entry is exhausted, the failed conversation is persisted to disk
+and a :class:`CascadeError` is raised.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from llm_pycascade.config import (
+    AppConfig,
+    CascadeConfig,
+    ProviderConfig,
+    ProviderType,
+    expand_tilde,
+)
+from llm_pycascade.db import is_on_cooldown, log_attempt, set_cooldown
+from llm_pycascade.error import CascadeError, ProviderError
+from llm_pycascade.persistence import save_failed_conversation
+from llm_pycascade.providers.anthropic import AnthropicProvider
+from llm_pycascade.providers.gemini import GeminiProvider
+from llm_pycascade.providers.ollama import OllamaProvider
+from llm_pycascade.providers.openai import OpenAIProvider
+
+if TYPE_CHECKING:
+    import aiosqlite
+
+    from llm_pycascade.models import Conversation, LlmResponse
+    from llm_pycascade.providers.base import LlmProvider
+
+logger = logging.getLogger(__name__)
+
+# ── constants ────────────────────────────────────────────────────────────
+
+BASE_COOLDOWN_SECS: float = 30.0
+MAX_COOLDOWN_SECS: float = 3600.0
+
+
+# ── provider factory ────────────────────────────────────────────────────
+
+
+def build_provider(
+    provider_name: str,
+    provider_config: ProviderConfig,
+    model: str,
+) -> LlmProvider:
+    """Instantiate the correct provider class based on the configuration.
+
+    Args:
+        provider_name: Logical name of the provider (for error messages).
+        provider_config: The provider's configuration block.
+        model: Model identifier to use.
+
+    Returns:
+        A fully-constructed :class:`LlmProvider` instance.
+
+    Raises:
+        ProviderError.missing_api_key: If no API key can be resolved.
+        ValueError: If the provider type is unknown.
+    """
+    from llm_pycascade.secrets import resolve_api_key
+
+    base_url = provider_config.base_url
+    api_key_env = provider_config.api_key_env or f"{provider_name.upper()}_API_KEY"
+    api_key_service = provider_config.api_key_service or provider_name
+
+    ptype = provider_config.type
+
+    if ptype == ProviderType.OPENAI:
+        api_key = resolve_api_key(api_key_service, api_key_env)
+        return OpenAIProvider(api_key=api_key, model=model, base_url=base_url)
+
+    if ptype == ProviderType.ANTHROPIC:
+        api_key = resolve_api_key(api_key_service, api_key_env)
+        return AnthropicProvider(api_key=api_key, model=model, base_url=base_url)
+
+    if ptype == ProviderType.GEMINI:
+        api_key = resolve_api_key(api_key_service, api_key_env)
+        return GeminiProvider(api_key=api_key, model=model, base_url=base_url)
+
+    if ptype == ProviderType.OLLAMA:
+        return OllamaProvider(model=model, base_url=base_url)
+
+    raise ValueError(f"Unknown provider type: {ptype}")
+
+
+# ── cooldown computation ──────────────────────────────────────────────────
+
+
+async def compute_cooldown(
+    entry_key: str,
+    conn: aiosqlite.Connection,
+) -> timedelta:
+    """Compute the cooldown duration based on recent failure history.
+
+    Uses exponential backoff:
+
+    * 1st failure  → ``BASE_COOLDOWN_SECS``
+    * 2nd failure  → ``2 × BASE_COOLDOWN_SECS``
+    * 3rd failure  → ``4 × BASE_COOLDOWN_SECS``
+    * …capped at ``MAX_COOLDOWN_SECS``
+
+    Args:
+        entry_key: ``provider/model`` identifier.
+        conn: Open database connection for querying failure history.
+
+    Returns:
+        The computed cooldown duration.
+    """
+    # Count recent failures within the last hour for this entry
+    one_hour_ago = time.time() - 3600.0
+    cursor = await conn.execute(
+        "SELECT COUNT(*) FROM attempt_log "
+        "WHERE provider_model = ? AND http_status IS NOT NULL AND timestamp > ?",
+        (entry_key, datetime.fromtimestamp(one_hour_ago, tz=timezone.utc).isoformat()),
+    )
+    row = await cursor.fetchone()
+    failure_count = row[0] if row else 0
+
+    # Exponential backoff: 2^n * BASE_COOLDOWN_SECS, capped
+
+    cooldown_secs = min(
+        (2**failure_count) * BASE_COOLDOWN_SECS,
+        MAX_COOLDOWN_SECS,
+    )
+    return timedelta(seconds=cooldown_secs)
+
+
+# ── main cascade function ──────────────────────────────────────────────
+
+
+async def run_cascade(
+    cascade_name: str,
+    conversation: Conversation,
+    config: AppConfig,
+    conn: aiosqlite.Connection,
+) -> LlmResponse:
+    """Run the named cascade against *conversation*.
+
+    Iterates through the cascade's entries in order:
+
+    1. Check if the entry is on cooldown → skip if so.
+    2. Build the provider and send the request.
+    3. On success: log the attempt and return the response.
+    4. On failure: log the attempt, compute & set a cooldown, and continue
+       to the next entry.
+    5. If all entries are exhausted: persist the failed conversation to disk
+       and raise :class:`CascadeError`.
+
+    Args:
+        cascade_name: Name of the cascade (must exist in ``config.cascades``).
+        conversation: The conversation to send.
+        config: The full application configuration.
+        conn: Open database connection for logging and cooldown queries.
+
+    Returns:
+        The first successful :class:`LlmResponse`.
+
+    Raises:
+        CascadeError: If every provider in the cascade fails.
+        KeyError: If the cascade name is not found in the configuration.
+    """
+    if cascade_name not in config.cascades:
+        raise KeyError(f"Cascade '{cascade_name}' not found in configuration")
+
+    cascade_config: CascadeConfig = config.cascades[cascade_name]
+    errors: list[tuple[str, ProviderError]] = []
+
+    for entry in cascade_config.entries:
+        provider_model_key = f"{entry.provider}/{entry.model}"
+
+        # 1. Check cooldown
+        if await is_on_cooldown(conn, provider_model_key):
+            logger.info(
+                "Skipping %s — on cooldown",
+                provider_model_key,
+            )
+            continue
+
+        # 2. Build provider
+        provider_cfg = config.providers.get(entry.provider)
+        if provider_cfg is None:
+            logger.warning(
+                "Provider '%s' not in config, skipping entry %s",
+                entry.provider,
+                provider_model_key,
+            )
+            continue
+
+        try:
+            provider = build_provider(entry.provider, provider_cfg, entry.model)
+        except (ProviderError, ValueError) as exc:
+            logger.error(
+                "Failed to build provider for %s: %s",
+                provider_model_key,
+                exc,
+            )
+            errors.append(
+                (
+                    provider_model_key,
+                    exc
+                    if isinstance(exc, ProviderError)
+                    else ProviderError.other(str(exc)),
+                )
+            )
+            continue
+
+        # 3. Send request
+        start = time.monotonic()
+        try:
+            response = await provider.complete(conversation)
+        except ProviderError as exc:
+            latency_ms = int((time.monotonic() - start) * 1000)
+
+            # 4. Log failure and set cooldown
+            logger.warning(
+                "Provider %s failed: %s",
+                provider_model_key,
+                exc,
+            )
+            await log_attempt(
+                conn,
+                cascade_name=cascade_name,
+                provider_model=provider_model_key,
+                http_status=exc.http_status,
+                latency_ms=latency_ms,
+                input_tokens=0,
+                output_tokens=0,
+            )
+
+            # Compute cooldown (respect Retry-After if available)
+            if exc.retry_after_seconds is not None:
+                cooldown_duration = timedelta(seconds=exc.retry_after_seconds)
+            else:
+                cooldown_duration = await compute_cooldown(provider_model_key, conn)
+
+            cooldown_until = time.time() + cooldown_duration.total_seconds()
+            await set_cooldown(conn, provider_model_key, cooldown_until)
+
+            logger.info(
+                "Cooldown set for %s until %s",
+                provider_model_key,
+                datetime.fromtimestamp(cooldown_until, tz=timezone.utc).isoformat(),
+            )
+
+            errors.append((provider_model_key, exc))
+            continue
+
+        latency_ms = int((time.monotonic() - start) * 1000)
+
+        # 5. Success — log and return
+        await log_attempt(
+            conn,
+            cascade_name=cascade_name,
+            provider_model=provider_model_key,
+            http_status=None,
+            latency_ms=latency_ms,
+            input_tokens=response.input_tokens,
+            output_tokens=response.output_tokens,
+        )
+        logger.info(
+            "Cascade '%s' succeeded via %s (%dms)",
+            cascade_name,
+            provider_model_key,
+            latency_ms,
+        )
+        return response
+
+    # 6. All entries exhausted — persist and raise
+    failure_dir = expand_tilde(config.failure_persistence.dir)
+    failed_path = save_failed_conversation(conversation, failure_dir, cascade_name)
+
+    error_summary = "; ".join(f"{k}: {v}" for k, v in errors)
+    raise CascadeError(
+        cascade_name=cascade_name,
+        message=f"All providers failed. {error_summary}",
+        failed_prompt_path=failed_path,
+    )

llm_pycascade/config.py

@@ -0,0 +1,193 @@
+"""Configuration types and TOML loading utilities."""
+
+from __future__ import annotations
+
+import os
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+# Python ≥ 3.11 has tomllib in the stdlib; fall back to tomli on older versions.
+if os.sys.version_info >= (3, 11):
+    import tomllib  # type: ignore[import-not-found]
+else:
+    try:
+        import tomllib  # type: ignore[import-not-found]
+    except ImportError:
+        import tomli as tomllib  # type: ignore[import-not-found]
+
+
+class ProviderType(str, Enum):
+    """Supported LLM provider types."""
+
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    GEMINI = "gemini"
+    OLLAMA = "ollama"
+
+
+class ProviderConfig(BaseModel):
+    """Configuration for a single LLM provider.
+
+    Attributes:
+        type: The provider type (openai, anthropic, gemini, ollama).
+        api_key_service: Service name to look up in the system keyring.
+        api_key_env: Environment variable name to read the API key from.
+        base_url: Override the default API base URL.
+    """
+
+    type: ProviderType
+    api_key_service: str | None = Field(default=None, exclude_none=True)
+    api_key_env: str | None = Field(default=None, exclude_none=True)
+    base_url: str | None = Field(default=None, exclude_none=True)
+
+
+class CascadeEntry(BaseModel):
+    """A single entry in a cascade — one provider/model pair.
+
+    Attributes:
+        provider: Key name of the provider (must match a key in ``[providers]``).
+        model: The model identifier to use.
+    """
+
+    provider: str
+    model: str
+
+
+class CascadeConfig(BaseModel):
+    """Ordered list of provider/model entries that form a cascade.
+
+    Attributes:
+        entries: The cascade entries, tried in order.
+    """
+
+    entries: list[CascadeEntry] = Field(default_factory=list)
+
+
+class DatabaseConfig(BaseModel):
+    """Configuration for the SQLite database used for attempt logging and cooldowns.
+
+    Attributes:
+        path: Filesystem path to the SQLite database.
+    """
+
+    path: str = "~/.local/share/llm-pycascade/db.sqlite"
+
+
+class FailureConfig(BaseModel):
+    """Configuration for persisting failed conversations.
+
+    Attributes:
+        dir: Directory where failed-prompt JSON files are saved.
+    """
+
+    dir: str = "~/.local/share/llm-pycascade/failed_prompts"
+
+
+class AppConfig(BaseModel):
+    """Top-level application configuration loaded from TOML.
+
+    Attributes:
+        providers: Mapping of provider names to their configurations.
+        cascades: Mapping of cascade names to ordered entry lists.
+        database: Database configuration.
+        failure_persistence: Failed-prompt persistence configuration.
+    """
+
+    providers: dict[str, ProviderConfig] = Field(default_factory=dict)
+    cascades: dict[str, CascadeConfig] = Field(default_factory=dict)
+    database: DatabaseConfig = Field(default_factory=DatabaseConfig)
+    failure_persistence: FailureConfig = Field(default_factory=FailureConfig)
+
+
+def expand_tilde(path: str) -> str:
+    """Expand a leading ``~`` to the user's home directory.
+
+    Args:
+        path: A filesystem path that may start with ``~``.
+
+    Returns:
+        The expanded absolute path.
+    """
+    return os.path.expanduser(path)
+
+
+def default_config_path() -> Path:
+    """Return the default configuration file path.
+
+    Searches (in order):
+
+    1. ``LLM_PYCASCADE_CONFIG`` environment variable
+    2. ``~/.config/llm-pycascade/config.toml``
+    3. ``~/.llm-pycascade.toml`` (legacy)
+
+    Returns:
+        Path to the first config file found, or the preferred default.
+    """
+    env_path = os.environ.get("LLM_PYCASCADE_CONFIG")
+    if env_path:
+        return Path(env_path)
+
+    xdg = Path.home() / ".config" / "llm-pycascade" / "config.toml"
+    if xdg.exists():
+        return xdg
+
+    legacy = Path.home() / ".llm-pycascade.toml"
+    if legacy.exists():
+        return legacy
+
+    return xdg
+
+
+def load_config(path: str | Path | None = None) -> AppConfig:
+    """Load and parse the application configuration from a TOML file.
+
+    Args:
+        path: Explicit path to the configuration file.  If ``None``, uses
+              :func:`default_config_path`.
+
+    Returns:
+        A fully-resolved :class:`AppConfig` instance.
+
+    Raises:
+        FileNotFoundError: If the config file does not exist.
+        tomllib.TOMLDecodeError: If the file contains invalid TOML.
+    """
+    path = default_config_path() if path is None else Path(path)
+
+    if not path.exists():
+        raise FileNotFoundError(f"Configuration file not found: {path}")
+
+    with open(path, "rb") as f:
+        raw: dict[str, Any] = tomllib.load(f)
+
+    # Parse the [providers] section
+    providers: dict[str, ProviderConfig] = {}
+    for name, pcfg in raw.get("providers", {}).items():
+        providers[name] = ProviderConfig(**pcfg)
+
+    # Parse the [cascades] section
+    cascades: dict[str, CascadeConfig] = {}
+    for cascade_name, ccfg in raw.get("cascades", {}).items():
+        entries_raw = ccfg.get("entries", [])
+        entries = [CascadeEntry(**e) for e in entries_raw]
+        cascades[cascade_name] = CascadeConfig(entries=entries)
+
+    # Parse optional [database] and [failure_persistence] sections
+    database = (
+        DatabaseConfig(**raw["database"]) if "database" in raw else DatabaseConfig()
+    )
+    failure = (
+        FailureConfig(**raw["failure_persistence"])
+        if "failure_persistence" in raw
+        else FailureConfig()
+    )
+
+    return AppConfig(
+        providers=providers,
+        cascades=cascades,
+        database=database,
+        failure_persistence=failure,
+    )

llm_pycascade/db.py

@@ -0,0 +1,123 @@
+"""Async SQLite database for attempt logging and cooldown tracking."""
+
+from __future__ import annotations
+
+import time
+
+import aiosqlite
+
+_CREATE_ATTEMPT_LOG = """
+CREATE TABLE IF NOT EXISTS attempt_log (
+    id              INTEGER PRIMARY KEY AUTOINCREMENT,
+    timestamp       TEXT NOT NULL DEFAULT (datetime('now')),
+    cascade_name    TEXT NOT NULL,
+    provider_model  TEXT NOT NULL,
+    http_status     INTEGER,
+    latency_ms      INTEGER,
+    input_tokens    INTEGER,
+    output_tokens   INTEGER
+);
+"""
+
+_CREATE_COOLDOWN = """
+CREATE TABLE IF NOT EXISTS cooldown (
+    provider_model  TEXT PRIMARY KEY,
+    cooldown_until  REAL NOT NULL
+);
+"""
+
+
+async def init_db(path: str) -> aiosqlite.Connection:
+    """Open (or create) the SQLite database and ensure tables exist.
+
+    Args:
+        path: Filesystem path to the SQLite database file.
+
+    Returns:
+        An open :class:`aiosqlite.Connection`.
+    """
+    conn = await aiosqlite.connect(path)
+    await conn.execute(_CREATE_ATTEMPT_LOG)
+    await conn.execute(_CREATE_COOLDOWN)
+    await conn.commit()
+    return conn
+
+
+async def log_attempt(
+    conn: aiosqlite.Connection,
+    cascade_name: str,
+    provider_model: str,
+    http_status: int | None,
+    latency_ms: int,
+    input_tokens: int,
+    output_tokens: int,
+) -> None:
+    """Insert a row into the attempt log.
+
+    Args:
+        conn: Open database connection.
+        cascade_name: Name of the cascade that was running.
+        provider_model: ``provider/model`` identifier string.
+        http_status: HTTP status code (``None`` if not HTTP-related).
+        latency_ms: Round-trip latency in milliseconds.
+        input_tokens: Input tokens reported by the provider.
+        output_tokens: Output tokens reported by the provider.
+    """
+    await conn.execute(
+        "INSERT INTO attempt_log "
+        "(cascade_name, provider_model, http_status, latency_ms, "
+        "input_tokens, output_tokens) "
+        "VALUES (?, ?, ?, ?, ?, ?)",
+        (
+            cascade_name,
+            provider_model,
+            http_status,
+            latency_ms,
+            input_tokens,
+            output_tokens,
+        ),
+    )
+    await conn.commit()
+
+
+async def is_on_cooldown(conn: aiosqlite.Connection, provider_model: str) -> bool:
+    """Check whether a provider/model combination is currently in cooldown.
+
+    Args:
+        conn: Open database connection.
+        provider_model: ``provider/model`` identifier string.
+
+    Returns:
+        ``True`` if the entry has an active cooldown.
+    """
+    cursor = await conn.execute(
+        "SELECT cooldown_until FROM cooldown WHERE provider_model = ?",
+        (provider_model,),
+    )
+    row = await cursor.fetchone()
+    if row is None:
+        return False
+    cooldown_until = row[0]
+    return time.time() < cooldown_until
+
+
+async def set_cooldown(
+    conn: aiosqlite.Connection,
+    provider_model: str,
+    cooldown_until: float,
+) -> None:
+    """Set (or update) the cooldown for a provider/model combination.
+
+    Uses ``INSERT OR REPLACE`` so repeated cooldowns simply update the expiry.
+
+    Args:
+        conn: Open database connection.
+        provider_model: ``provider/model`` identifier string.
+        cooldown_until: Unix timestamp (seconds since epoch) when the cooldown expires.
+    """
+    await conn.execute(
+        "INSERT OR REPLACE INTO cooldown "
+        "(provider_model, cooldown_until) VALUES (?, ?)",
+        (provider_model, cooldown_until),
+    )
+    await conn.commit()