Multi-Provider AI Load Balancer

Problem Statement

Relying on a single AI provider for production infrastructure introduces critical vulnerabilities: strict rate limits, vendor lock-in, unexpected pricing adjustments, and API downtime. When an upstream provider experiences service degradation or rate limiting (HTTP 429), dependent applications crash or stall, degrading user experience and breaking SLA commitments.

A Multi-Provider AI Load Balancer resolves these issues by abstracts provider-specific SDKs behind a unified interface. It dynamically distributes prompt requests across multiple Large Language Model (LLM) vendors—such as OpenAI, DeepSeek, xAI (Grok), and Google Gemini. Featuring automated retries, rate-limit exponential backoffs, and seamless provider failover, this architecture ensures high availability, reduced latency, and resilient LLM operations.

How It Works

The load balancer operates on a unified abstraction layer. Each supported LLM provider inherits from a standardized interface (BaseAIProvider) that encapsulates authentication, API client instantiation, and request formatting. Provider instances normalize output formats and timing metrics into a unified ProviderResponse dataclass.

When a prompt is submitted:

1. Provider Resolution: The script instantiates active provider drivers based on command-line flags and verified environment variables.
2. Execution Strategy: Under the failover strategy, the system attempts generation via the primary provider. If transient errors occur (e.g., rate limits, network timeouts), tenacity executes local exponential backoff retries. If local retries exhaust or fail critically, the balancer automatically falls back to the next available provider in the pipeline.
3. Round-Robin Routing: Under round-robin, the balancer rotates the starting provider index sequentially across executions, balancing load evenly across configured providers before executing fallback logic.

                  +-------------------------+
                  |  Load Balancer Pipeline |
                  +------------+------------+
                               |
            +------------------+------------------+
            | (Failover / Round-Robin Scheduling) |
            +------------------+------------------+
                               |
      +------------------------+------------------------+
      |                        |                        |
+-----v-------+          +-----v-------+          +-----v-------+
|  DeepSeek   |          |   Gemini    |          |    Grok     |
| (OpenAI SDK)|          | (GenAI SDK) |          | (OpenAI SDK)|
+-------------+          +-------------+          +-------------+

Full Script

#!/usr/bin/env python3
"""
Multi-Provider AI Load Balancer
A resilient interface to route, retry, and failover prompts across DeepSeek,
Gemini, Grok, and OpenAI APIs.
"""

import argparse
import dataclasses
import logging
import os
import sys
import time
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Type

from google import genai
from google.genai import types as genai_types
from openai import OpenAI, OpenAIError
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("ai_balancer")


@dataclasses.dataclass
class ProviderResponse:
    """Standardized response container across all providers."""
    content: str
    provider: str
    model: str
    latency_ms: float


class BaseAIProvider(ABC):
    """Abstract Base Class for AI Provider Implementations."""

    def __init__(self, api_key: str, model_name: str):
        self.api_key = api_key
        self.model_name = model_name

    @abstractmethod
    def generate(self, prompt: str, max_tokens: int = 500) -> str:
        """Execute the prompt request against the underlying provider API."""
        pass


class OpenAICompatibleProvider(BaseAIProvider):
    """Handles OpenAI-compatible APIs (OpenAI, DeepSeek, Grok)."""

    def __init__(self, api_key: str, model_name: str, base_url: Optional[str] = None):
        super().__init__(api_key, model_name)
        self.client = OpenAI(api_key=self.api_key, base_url=base_url)

    def generate(self, prompt: str, max_tokens: int = 500) -> str:
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=max_tokens,
        )
        return response.choices[0].message.content or ""


class GeminiProvider(BaseAIProvider):
    """Handles Google GenAI SDK integration."""

    def __init__(self, api_key: str, model_name: str):
        super().__init__(api_key, model_name)
        self.client = genai.Client(api_key=self.api_key)

    def generate(self, prompt: str, max_tokens: int = 500) -> str:
        config = genai_types.GenerateContentConfig(max_output_tokens=max_tokens)
        response = self.client.models.generate_content(
            model=self.model_name,
            contents=prompt,
            config=config,
        )
        return response.text or ""


class LoadBalancer:
    """Manages provider routing, retries, and failovers."""

    DEFAULT_MODELS = {
        "openai": ("gpt-4o-mini", None),
        "deepseek": ("deepseek-chat", "https://api.deepseek.com"),
        "grok": ("grok-2-latest", "https://api.x.ai/v1"),
        "gemini": ("gemini-2.5-flash", None),
    }

    ENV_KEYS = {
        "openai": "OPENAI_API_KEY",
        "deepseek": "DEEPSEEK_API_KEY",
        "grok": "GROK_API_KEY",
        "gemini": "GEMINI_API_KEY",
    }

    def __init__(self, requested_providers: List[str]):
        self.providers: Dict[str, BaseAIProvider] = {}
        self.provider_order: List[str] = []
        self._rr_index = 0

        self._initialize_providers(requested_providers)

    def _initialize_providers(self, requested_providers: List[str]) -> None:
        """Instantiates drivers for requested providers with valid credentials."""
        for key in requested_providers:
            key = key.strip().lower()
            if key not in self.DEFAULT_MODELS:
                logger.warning(f"Unknown provider '{key}'. Skipping.")
                continue

            env_var = self.ENV_KEYS[key]
            api_key = os.environ.get(env_var)
            if not api_key:
                logger.warning(f"Environment variable '{env_var}' missing for provider '{key}'. Skipping.")
                continue

            model_name, base_url = self.DEFAULT_MODELS[key]
            if key == "gemini":
                provider_inst = GeminiProvider(api_key=api_key, model_name=model_name)
            else:
                provider_inst = OpenAICompatibleProvider(
                    api_key=api_key, model_name=model_name, base_url=base_url
                )

            self.providers[key] = provider_inst
            self.provider_order.append(key)

        if not self.providers:
            raise RuntimeError("No valid AI providers were initialized. Check API keys and provider names.")

    def _execute_with_retry(self, provider: BaseAIProvider, prompt: str, max_tokens: int) -> str:
        """Executes API call with exponential backoff via Tenacity."""
        @retry(
            reraise=True,
            stop=stop_after_attempt(3),
            wait=wait_exponential(multiplier=1, min=2, max=10),
            retry=retry_if_exception_type(Exception),
        )
        def _call():
            return provider.generate(prompt, max_tokens=max_tokens)

        return _call()

    def generate(self, prompt: str, strategy: str = "failover", max_tokens: int = 500) -> ProviderResponse:
        """Routes prompt through requested providers using selected load balancing strategy."""
        if strategy == "round-robin":
            # Rotate sequence based on current round-robin counter
            order = self.provider_order[self._rr_index:] + self.provider_order[:self._rr_index]
            self._rr_index = (self._rr_index + 1) % len(self.provider_order)
        else:
            order = list(self.provider_order)

        errors = []
        for name in order:
            provider = self.providers[name]
            logger.info(f"Attempting prompt generation using provider: '{name}' ({provider.model_name})")
            start_time = time.perf_counter()

            try:
                content = self._execute_with_retry(provider, prompt, max_tokens)
                elapsed_ms = (time.perf_counter() - start_time) * 1000
                logger.info(f"Successfully generated response via '{name}' in {elapsed_ms:.2f}ms")
                return ProviderResponse(
                    content=content,
                    provider=name,
                    model=provider.model_name,
                    latency_ms=elapsed_ms,
                )
            except Exception as err:
                logger.warning(f"Provider '{name}' failed after retries: {err}. Failing over...")
                errors.append(f"{name}: {str(err)}")

        raise RuntimeError(f"All providers exhausted. Failure trace: {'; '.join(errors)}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Multi-Provider AI Load Balancer")
    parser.add_argument(
        "--providers",
        type=str,
        default="deepseek,gemini,grok,openai",
        help="Comma-separated list of providers in priority order (e.g., 'deepseek,gemini,grok')",
    )
    parser.add_argument(
        "--prompt",
        type=str,
        default="Explain the concept of circuit breakers in microservices in 2 concise sentences.",
        help="Prompt text to evaluate.",
    )
    parser.add_argument(
        "--strategy",
        type=str,
        choices=["failover", "round-robin"],
        default="failover",
        help="Load balancing strategy to execute.",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=250,
        help="Maximum tokens for output response generation.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    requested_providers = [p.strip() for p in args.providers.split(",") if p.strip()]

    try:
        balancer = LoadBalancer(requested_providers=requested_providers)
        response = balancer.generate(
            prompt=args.prompt, strategy=args.strategy, max_tokens=args.max_tokens
        )

        print("\n" + "=" * 60)
        print(f"PROVIDER : {response.provider.upper()} ({response.model})")
        print(f"LATENCY  : {response.latency_ms:.2f} ms")
        print("=" * 60)
        print(f"RESPONSE :\n{response.content.strip()}")
        print("=" * 60 + "\n")

    except Exception as exc:
        logger.error(f"Execution failed: {exc}")
        sys.exit(1)


if __name__ == "__main__":
    main()

Configuration

The script reads provider credentials from system environment variables. Set the appropriate key for every provider specified in your execution pipeline:

Export the necessary keys in your shell prior to execution:

export DEEPSEEK_API_KEY="sk-..."
export GEMINI_API_KEY="AIzaSy..."
export GROK_API_KEY="xai-..."
export OPENAI_API_KEY="sk-proj-..."

Usage Examples

Example 1: Standard Priority Failover

Request execution prioritizing DeepSeek, then Gemini, then Grok:

python3 ai_balancer.py --providers deepseek,gemini,grok --prompt "Define idempotency in microservices."

Expected Output:

2026-03-30 10:15:00,102 [INFO] Attempting prompt generation using provider: 'deepseek' (deepseek-chat)
2026-03-30 10:15:01,245 [INFO] Successfully generated response via 'deepseek' in 1143.20ms

============================================================
PROVIDER : DEEPSEEK (deepseek-chat)
LATENCY  : 1143.20 ms
============================================================
RESPONSE :
Idempotency means that making multiple identical requests has the same effect as making a single request. In microservices, it ensures that operations like retried API calls or duplicate message deliveries do not alter system state unexpectedly.
============================================================

Example 2: Automatic Provider Failover on Rate Limit / Missing Key

If the primary provider (deepseek) experiences an outage or has exhausted its quota, the load balancer catches the error, triggers backoff retries, and seamlessly falls back to the secondary provider (gemini):

python3 ai_balancer.py --providers deepseek,gemini --strategy failover

Expected Output:

2026-03-30 10:16:00,050 [INFO] Attempting prompt generation using provider: 'deepseek' (deepseek-chat)
2026-03-30 10:16:02,110 [WARNING] Provider 'deepseek' failed after retries: Error code: 429 - Rate limit reached. Failing over...
2026-03-30 10:16:02,111 [INFO] Attempting prompt generation using provider: 'gemini' (gemini-2.5-flash)
2026-03-30 10:16:02,890 [INFO] Successfully generated response via 'gemini' in 779.12ms

============================================================
PROVIDER : GEMINI (gemini-2.5-flash)
LATENCY  : 779.12 ms
============================================================
RESPONSE :
A circuit breaker monitors remote service calls and immediately fails requests when the target service is detected to be unhealthy, preventing cascading system failure and resource exhaustion across microservices.
============================================================

Error Handling

1. Missing Provider Keys: On initialization, the script scans environment variables for each provider requested via --providers. If a key is missing, a warning is logged, and the missing provider is gracefully omitted from the active pipeline without crashing.
2. Per-Provider Retries (Tenacity): Network hiccups, temporary server errors (HTTP 5xx), or rate limits (HTTP 429) trigger a 3-attempt exponential backoff wait (between 2 to 10 seconds) before declaring a provider attempt failed.
3. Total Pipeline Exhaustion: If every provider in the configured pipeline fails local retries, the script raises a consolidated RuntimeError listing all accumulated failure reasons across every attempted backend, terminating with exit code 1.

Extending Ideas

• Dynamic Health Tracking & Cooldowns: Implement a circuit breaker state machine (e.g., mark a provider as "unhealthy" for 5 minutes after consecutive failures) to skip degraded endpoints entirely rather than waiting for timeouts.
• Latency & Cost Optimization Routing: Track dynamic request latency and token pricing. Dynamically route lower-priority prompts to cheaper/faster models (e.g., DeepSeek/Gemini Flash) and reserve premium models (GPT-4o) for high-complexity requests.
• Streaming Response Support: Extend BaseAIProvider and LoadBalancer to yield server-sent event (SSE) token streams, delivering real-time streaming text to frontend interfaces regardless of which backend is actively serving the query.