A reusable LLM inference module that connects to Ollama, OpenAI, or OpenRouter. It selects a preferred provider through .env and automatically falls back to another provider when the preferred service is unavailable or not configured.

llm_provider_fallback.py

"""LLM inference with automatic Ollama, OpenAI, and OpenRouter fallback.

Setup:
1. Install packages: py -m pip install -r requirements.txt
2. Create a .env file in this project's root directory.
3. Set AI_PROVIDER to the provider that should be tried first:
   AI_PROVIDER=ollama        # Local inference; no API key required.
   AI_PROVIDER=openai        # Requires OPENAI_API_KEY.
   AI_PROVIDER=openrouter    # Requires OPENROUTER_API_KEY.

Add the variables needed by the providers you want to use:

    AI_PROVIDER=openrouter
    OPENAI_API_KEY=your-openai-key
    OPENAI_MODEL=gpt-4o-mini
    OPENROUTER_API_KEY=your-openrouter-key
    OPENROUTER_MODEL=openai/gpt-4o-mini
    OLLAMA_BASE_URL=http://localhost:11434/v1
    OLLAMA_MODEL=llama3.2

For Ollama, install and start Ollama and pull the configured model first:
    ollama pull llama3.2

Run this file with: py llm_provider_fallback.py

The selected provider is tried first. Missing keys and failed requests cause the
program to try the remaining configured providers automatically.
"""

import os
import sys
from dataclasses import dataclass

from dotenv import load_dotenv
from openai import OpenAI, OpenAIError
from openai.types.chat import ChatCompletionMessageParam

load_dotenv()

PROVIDER_NAMES = ("ollama", "openai", "openrouter")
PLACEHOLDER_KEYS = {
    "sk-your-key-here",
    "sk-your-openai-key-here",
    "sk-your-openrouter-key-here",
}


@dataclass(frozen=True)
class ProviderConfig:
    """Connection settings for one OpenAI-compatible provider."""

    name: str
    base_url: str
    api_key: str
    model: str


@dataclass(frozen=True)
class InferenceResult:
    """Generated text and the provider that returned it."""

    provider: str
    content: str


def cloud_api_key(variable: str) -> str:
    """Return a configured cloud key or report that provider as unavailable."""
    api_key = os.getenv(variable)
    if not api_key or api_key in PLACEHOLDER_KEYS:
        raise RuntimeError(f"{variable} is not configured")
    return api_key


def provider_config(name: str) -> ProviderConfig:
    """Build connection settings from environment variables."""
    if name == "ollama":
        return ProviderConfig(
            name="ollama",
            base_url=os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1"),
            api_key="ollama",
            model=os.getenv("OLLAMA_MODEL", "llama3.2"),
        )
    if name == "openai":
        return ProviderConfig(
            name="openai",
            base_url="https://api.openai.com/v1",
            api_key=cloud_api_key("OPENAI_API_KEY"),
            model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
        )
    if name == "openrouter":
        return ProviderConfig(
            name="openrouter",
            base_url="https://openrouter.ai/api/v1",
            api_key=cloud_api_key("OPENROUTER_API_KEY"),
            model=os.getenv("OPENROUTER_MODEL", "openai/gpt-4o-mini"),
        )
    raise RuntimeError(f"Unsupported provider: {name}")


def provider_order() -> tuple[str, ...]:
    """Put the selected provider first while retaining automatic fallbacks."""
    preferred = os.getenv("AI_PROVIDER", "openrouter").strip().lower()
    if preferred not in PROVIDER_NAMES:
        choices = ", ".join(PROVIDER_NAMES)
        raise RuntimeError(f"AI_PROVIDER must be one of: {choices}")
    return (preferred, *(name for name in PROVIDER_NAMES if name != preferred))


def call_provider(
    provider: ProviderConfig,
    messages: list[ChatCompletionMessageParam],
) -> InferenceResult:
    """Send one inference request to a configured provider."""
    client = OpenAI(
        api_key=provider.api_key,
        base_url=provider.base_url,
        max_retries=0,
        timeout=30.0,
    )
    response = client.chat.completions.create(
        model=provider.model,
        messages=messages,
        temperature=0.3,
        max_tokens=300,
    )
    return InferenceResult(
        provider=provider.name,
        content=response.choices[0].message.content or "",
    )


def infer(messages: list[ChatCompletionMessageParam]) -> InferenceResult:
    """Try each provider until one returns an inference result."""
    failures: list[str] = []
    for name in provider_order():
        try:
            provider = provider_config(name)
            return call_provider(provider, messages)
        except (OpenAIError, RuntimeError) as error:
            failures.append(f"{name}: {type(error).__name__}")
            print(f"{name} unavailable; trying next provider.", file=sys.stderr)

    failure_summary = ", ".join(failures)
    raise RuntimeError(f"No LLM provider succeeded ({failure_summary}).")


def main() -> None:
    """Run a small inference example using provider fallback."""
    messages: list[ChatCompletionMessageParam] = [
        {
            "role": "system",
            "content": "You are a helpful AfyaPlus health assistant.",
        },
        {
            "role": "user",
            "content": "Say hello in Swahili.",
        },
    ]
    try:
        result = infer(messages)
        print(f"Provider: {result.provider}")
        print(result.content)
    except RuntimeError as error:
        print(f"Inference failed: {error}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()