fallback.py

from __future__ import annotations

import logging
import os
import random
import time
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime
from typing import Any

from openai import (
    APIConnectionError,
    APIStatusError,
    APITimeoutError,
    OpenAI,
)


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
logger = logging.getLogger(__name__)


# These errors normally require a configuration or request change.
# Retrying the same request is unlikely to fix them.
FAIL_FAST_STATUS_CODES = {
    400,  # Invalid request or unsupported parameter
    401,  # Invalid or missing credential
    403,  # Permission denied
    404,  # Endpoint or model not found
    422,  # Request could not be processed
}


# These errors may be temporary.
TRANSIENT_STATUS_CODES = {
    408,  # Request timeout
    429,  # Rate limit
    500,  # Internal server error
    502,  # Bad gateway
    503,  # Service unavailable
    504,  # Gateway timeout
}


ATTEMPTS_PER_MODEL = 3
BACKOFF_BASE_SECONDS = 1.0
BACKOFF_CAP_SECONDS = 20.0
REQUEST_TIMEOUT_SECONDS = 30.0


class InvalidModelResponse(RuntimeError):
    """Raised when a successful HTTP response is unusable."""


def load_config() -> tuple[str, str, tuple[str, ...]]:
    """Read and validate environment variables."""

    api_key = os.getenv("COMETAPI_API_KEY", "").strip()

    base_url = os.getenv(
        "COMETAPI_BASE_URL",
        "https://api.cometapi.com/v1",
    ).strip()

    models = tuple(
        model.strip()
        for model in os.getenv(
            "COMETAPI_MODEL_CHAIN",
            "",
        ).split(",")
        if model.strip()
    )

    if not api_key:
        raise RuntimeError(
            "COMETAPI_API_KEY is missing. "
            "Set it as an environment variable."
        )

    if not models:
        raise RuntimeError(
            "COMETAPI_MODEL_CHAIN is missing. "
            "Provide one or more comma-separated model IDs."
        )

    return api_key, base_url, models


def parse_retry_after(value: str | None) -> float | None:
    """
    Parse Retry-After as either a number of seconds or an HTTP date.

    Return None when the header is missing or invalid.
    """

    if not value:
        return None

    value = value.strip()

    try:
        return max(0.0, float(value))
    except ValueError:
        pass

    try:
        retry_time = parsedate_to_datetime(value)

        if retry_time.tzinfo is None:
            retry_time = retry_time.replace(
                tzinfo=timezone.utc,
            )

        return max(
            0.0,
            (
                retry_time
                - datetime.now(timezone.utc)
            ).total_seconds(),
        )

    except (
        TypeError,
        ValueError,
        OverflowError,
    ):
        return None


def calculate_delay(
    attempt_index: int,
    retry_after: float | None = None,
) -> float:
    """
    Calculate the next retry delay.

    Retry-After takes precedence when the server provides it.
    Otherwise, use exponential backoff with full jitter.
    """

    if retry_after is not None:
        return retry_after

    upper_bound = min(
        BACKOFF_CAP_SECONDS,
        BACKOFF_BASE_SECONDS * (2**attempt_index),
    )

    return random.uniform(
        0.0,
        upper_bound,
    )


def extract_text(response: Any) -> str:
    """
    Validate and extract text from a Chat Completions response.

    HTTP success alone does not guarantee a usable result.
    """

    choices = getattr(
        response,
        "choices",
        None,
    )

    if not choices:
        raise InvalidModelResponse(
            "The response contains no choices."
        )

    message = getattr(
        choices[0],
        "message",
        None,
    )

    content = (
        getattr(message, "content", None)
        if message
        else None
    )

    if not isinstance(content, str):
        raise InvalidModelResponse(
            "The response does not contain text."
        )

    if not content.strip():
        raise InvalidModelResponse(
            "The response contains empty text."
        )

    return content.strip()


def generate_with_fallback(
    client: OpenAI,
    models: tuple[str, ...],
    messages: list[dict[str, str]],
) -> str:
    """
    Retry temporary failures, then move to the next model.

    Fail-fast errors stop the entire operation.
    """

    failures: list[str] = []

    for model in models:
        for attempt_index in range(
            ATTEMPTS_PER_MODEL
        ):
            attempt_number = attempt_index + 1
            started_at = time.perf_counter()

            try:
                response = (
                    client.chat.completions.create(
                        model=model,
                        messages=messages,
                    )
                )

                text = extract_text(response)

                latency_ms = round(
                    (
                        time.perf_counter()
                        - started_at
                    )
                    * 1000
                )

                logger.info(
                    "request_succeeded "
                    "model=%s "
                    "attempt=%d "
                    "latency_ms=%d",
                    model,
                    attempt_number,
                    latency_ms,
                )

                return text

            except APIStatusError as exc:
                status_code = exc.status_code

                headers = getattr(
                    exc.response,
                    "headers",
                    {},
                )

                request_id = headers.get(
                    "x-request-id",
                    "-",
                )

                logger.warning(
                    "request_failed "
                    "model=%s "
                    "attempt=%d "
                    "status=%s "
                    "request_id=%s",
                    model,
                    attempt_number,
                    status_code,
                    request_id,
                )

                if (
                    status_code
                    in FAIL_FAST_STATUS_CODES
                ):
                    raise RuntimeError(
                        "Non-retryable HTTP error: "
                        f"model={model} "
                        f"status={status_code} "
                        f"request_id={request_id}"
                    ) from exc

                if (
                    status_code
                    not in TRANSIENT_STATUS_CODES
                ):
                    raise RuntimeError(
                        "Unclassified HTTP error: "
                        f"model={model} "
                        f"status={status_code} "
                        f"request_id={request_id}"
                    ) from exc

                failures.append(
                    f"{model}:HTTP_{status_code}"
                )

                if (
                    attempt_number
                    < ATTEMPTS_PER_MODEL
                ):
                    retry_after = (
                        parse_retry_after(
                            headers.get(
                                "retry-after"
                            )
                        )
                    )

                    delay = calculate_delay(
                        attempt_index,
                        retry_after,
                    )

                    logger.info(
                        "retry_scheduled "
                        "model=%s "
                        "delay_seconds=%.2f",
                        model,
                        delay,
                    )

                    time.sleep(delay)
                    continue

                break

            except (
                APITimeoutError,
                APIConnectionError,
            ) as exc:
                error_type = type(exc).__name__

                failures.append(
                    f"{model}:{error_type}"
                )

                logger.warning(
                    "transport_error "
                    "model=%s "
                    "attempt=%d "
                    "error_type=%s",
                    model,
                    attempt_number,
                    error_type,
                )

                if (
                    attempt_number
                    < ATTEMPTS_PER_MODEL
                ):
                    delay = calculate_delay(
                        attempt_index
                    )

                    logger.info(
                        "retry_scheduled "
                        "model=%s "
                        "delay_seconds=%.2f",
                        model,
                        delay,
                    )

                    time.sleep(delay)
                    continue

                break

            except InvalidModelResponse as exc:
                failures.append(
                    f"{model}:INVALID_RESPONSE"
                )

                logger.warning(
                    "invalid_response "
                    "model=%s "
                    "attempt=%d "
                    "reason=%s",
                    model,
                    attempt_number,
                    exc,
                )

                if (
                    attempt_number
                    < ATTEMPTS_PER_MODEL
                ):
                    delay = calculate_delay(
                        attempt_index
                    )

                    time.sleep(delay)
                    continue

                break

        logger.warning(
            "model_attempts_exhausted "
            "model=%s "
            "moving_to_next_model=true",
            model,
        )

    failure_summary = (
        ", ".join(failures)
        if failures
        else "none"
    )

    raise RuntimeError(
        "Every configured model failed. "
        f"Failure summary: {failure_summary}"
    )


def main() -> None:
    api_key, base_url, models = load_config()

    client = OpenAI(
        api_key=api_key,
        base_url=base_url,
        timeout=REQUEST_TIMEOUT_SECONDS,

        # Disable hidden SDK retries because this file
        # implements its own retry and fallback policy.
        max_retries=0,
    )

    messages = [
        {
            "role": "system",
            "content": (
                "Answer clearly. "
                "Do not invent information that is "
                "not supported by the prompt."
            ),
        },
        {
            "role": "user",
            "content": (
                "Explain the difference between "
                "retrying a request and falling back "
                "to another model."
            ),
        },
    ]

    result = generate_with_fallback(
        client=client,
        models=models,
        messages=messages,
    )

    print(result)


if __name__ == "__main__":
    main()