Module 8: Unified AI Client — Final integrating project

Base implementation

Time to type code. In this capsule you build the skeleton of the unified_ai_client package with the first adapter working (OpenAI). By the end you'll be able to do pip install -e . and use your library in any script.

By the end you'll be able to:

  • Create an installable Python structure with pyproject.toml
  • Implement BaseAdapter and OpenAIAdapter with full typing
  • Build a minimum viable UnifiedClient (only primary, no fallback yet)
  • Load configuration from YAML
  • Verify that everything works with a test script

Package setup

Create the file structure:

mkdir unified_ai_client_pkg
cd unified_ai_client_pkg
mkdir -p unified_ai_client/adapters tests
touch unified_ai_client/__init__.py
touch unified_ai_client/adapters/__init__.py
touch unified_ai_client/{models.py,exceptions.py,factory.py,client.py}
touch unified_ai_client/adapters/{base.py,openai_adapter.py}
touch tests/__init__.py
touch README.md pyproject.toml

Your tree:

unified_ai_client_pkg/
├── pyproject.toml
├── README.md
├── unified_ai_client/
│   ├── __init__.py
│   ├── models.py
│   ├── exceptions.py
│   ├── factory.py
│   ├── client.py
│   └── adapters/
│       ├── __init__.py
│       ├── base.py
│       └── openai_adapter.py
└── tests/
    └── __init__.py

pyproject.toml

[build-system]
requires = ["setuptools>=68.0"]
build-backend = "setuptools.build_meta"

[project]
name = "unified-ai-client"
version = "0.1.0"
description = "Unified client for multiple LLM providers"
requires-python = ">=3.10"
dependencies = [
    "openai>=1.30.0",
    "pydantic>=2.0",
    "pyyaml>=6.0",
    "httpx>=0.25",
]

[project.optional-dependencies]
dev = ["pytest", "pytest-mock"]

[tool.setuptools.packages.find]
include = ["unified_ai_client*"]

Install:

python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"

pip list should show unified-ai-client 0.1.0.


unified_ai_client/models.py

# unified_ai_client/models.py
from typing import Literal
from pydantic import BaseModel, Field


class Message(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str


class ChatResponse(BaseModel):
    text: str
    model: str
    provider: str
    tokens_input: int = 0
    tokens_output: int = 0
    duration_ms: int = 0
    cost_usd: float | None = None
    raw_response: dict | None = None


class ProviderConfig(BaseModel):
    name: str
    type: Literal["openai", "openai_compatible", "modal_custom"]
    model: str | None = None
    base_url: str | None = None
    api_key_env: str | None = None
    base_url_env: str | None = None
    api_token_env: str | None = None
    price_input_per_1m: float | None = None
    price_output_per_1m: float | None = None


class ClientConfig(BaseModel):
    primary: str
    fallback: list[str] = Field(default_factory=list)
    providers: dict[str, ProviderConfig]
    metrics_enabled: bool = True

unified_ai_client/exceptions.py

# unified_ai_client/exceptions.py
class UnifiedClientError(Exception):
    """Base for client errors."""


class ConfigError(UnifiedClientError):
    """Configuration error (provider doesn't exist, missing env var, etc.)."""


class ProviderError(UnifiedClientError):
    """Error reported by a specific provider."""

    def __init__(self, provider: str, message: str, original: Exception | None = None):
        super().__init__(f"[{provider}] {message}")
        self.provider = provider
        self.original = original


class AuthError(ProviderError):
    """Invalid API key, expired token, etc."""


class RateLimitError(ProviderError):
    """Provider returned 429."""


class TimeoutError(ProviderError):
    """Request exceeded the timeout."""


class AllProvidersFailedError(UnifiedClientError):
    """Primary and all fallbacks failed."""

    def __init__(self, errors: dict[str, Exception]):
        msg = "All providers failed:\n" + "\n".join(
            f"  - {p}: {e}" for p, e in errors.items()
        )
        super().__init__(msg)
        self.errors = errors

unified_ai_client/adapters/base.py

# unified_ai_client/adapters/base.py
from abc import ABC, abstractmethod
from ..models import Message, ChatResponse, ProviderConfig


class BaseAdapter(ABC):
    def __init__(self, config: ProviderConfig):
        self.config = config
        self.name = config.name

    @abstractmethod
    def chat(
        self,
        messages: list[Message],
        max_tokens: int = 256,
        temperature: float = 0.7,
    ) -> ChatResponse:
        """Generates a response from messages. Per-provider implementation."""
        ...

unified_ai_client/adapters/openai_adapter.py

Covers native OpenAI and all OpenAI-compatible providers (OpenRouter, Ollama, LM Studio).

# unified_ai_client/adapters/openai_adapter.py
import os
import time
from openai import OpenAI, APIStatusError, RateLimitError as OpenAIRateLimitError
from openai import APIConnectionError, AuthenticationError, APITimeoutError

from .base import BaseAdapter
from ..models import Message, ChatResponse, ProviderConfig
from ..exceptions import (
    ProviderError,
    AuthError,
    RateLimitError,
    TimeoutError,
    ConfigError,
)


class OpenAIAdapter(BaseAdapter):
    """
    Works for native OpenAI and any OpenAI-compatible provider:
    OpenRouter, Ollama, LM Studio, Anyscale, Together, etc.
    """

    def __init__(self, config: ProviderConfig):
        super().__init__(config)
        api_key = self._resolve_api_key(config)
        base_url = config.base_url or self._resolve_base_url(config)
        self.client = OpenAI(api_key=api_key, base_url=base_url, timeout=60.0)

    def _resolve_api_key(self, config: ProviderConfig) -> str:
        if not config.api_key_env:
            return "no-key-needed"  # local Ollama case
        key = os.environ.get(config.api_key_env)
        if not key:
            raise ConfigError(
                f"Missing env var '{config.api_key_env}' for provider '{config.name}'"
            )
        return key

    def _resolve_base_url(self, config: ProviderConfig) -> str | None:
        if config.base_url_env:
            return os.environ.get(config.base_url_env)
        return None

    def chat(
        self,
        messages: list[Message],
        max_tokens: int = 256,
        temperature: float = 0.7,
    ) -> ChatResponse:
        if not self.config.model:
            raise ConfigError(f"Provider '{self.name}' has no 'model' configured")

        openai_msgs = [{"role": m.role, "content": m.content} for m in messages]

        start = time.perf_counter()
        try:
            r = self.client.chat.completions.create(
                model=self.config.model,
                messages=openai_msgs,
                max_tokens=max_tokens,
                temperature=temperature,
            )
        except AuthenticationError as e:
            raise AuthError(self.name, "Invalid API key or no permissions", e) from e
        except OpenAIRateLimitError as e:
            raise RateLimitError(self.name, "Rate limit reached", e) from e
        except APITimeoutError as e:
            raise TimeoutError(self.name, "Request timeout", e) from e
        except (APIConnectionError, APIStatusError) as e:
            raise ProviderError(self.name, f"Provider error: {e}", e) from e

        duration_ms = int((time.perf_counter() - start) * 1000)

        text = r.choices[0].message.content or ""
        usage = r.usage
        tokens_input = usage.prompt_tokens if usage else 0
        tokens_output = usage.completion_tokens if usage else 0

        cost = self._calculate_cost(tokens_input, tokens_output)

        return ChatResponse(
            text=text,
            model=self.config.model,
            provider=self.name,
            tokens_input=tokens_input,
            tokens_output=tokens_output,
            duration_ms=duration_ms,
            cost_usd=cost,
            raw_response=r.model_dump(),
        )

    def _calculate_cost(self, tokens_input: int, tokens_output: int) -> float | None:
        if self.config.price_input_per_1m is None or self.config.price_output_per_1m is None:
            return None
        return (
            tokens_input * self.config.price_input_per_1m / 1_000_000
            + tokens_output * self.config.price_output_per_1m / 1_000_000
        )

unified_ai_client/factory.py

# unified_ai_client/factory.py
from .models import ProviderConfig
from .adapters.base import BaseAdapter
from .adapters.openai_adapter import OpenAIAdapter
from .exceptions import ConfigError


class ProviderFactory:
    """Creates adapters given a ProviderConfig."""

    _REGISTRY: dict[str, type[BaseAdapter]] = {
        "openai": OpenAIAdapter,
        "openai_compatible": OpenAIAdapter,
        # "modal_custom" will be added in a later capsule
    }

    @classmethod
    def create(cls, config: ProviderConfig) -> BaseAdapter:
        adapter_class = cls._REGISTRY.get(config.type)
        if not adapter_class:
            raise ConfigError(
                f"Unknown provider type: '{config.type}'. "
                f"Valid types: {list(cls._REGISTRY.keys())}"
            )
        return adapter_class(config)

    @classmethod
    def register(cls, provider_type: str, adapter_class: type[BaseAdapter]) -> None:
        """Registers a custom adapter from outside (extensibility)."""
        cls._REGISTRY[provider_type] = adapter_class

unified_ai_client/client.py

# unified_ai_client/client.py
from pathlib import Path
from typing import Iterable
import yaml

from .models import Message, ChatResponse, ClientConfig
from .factory import ProviderFactory
from .exceptions import ConfigError


class UnifiedClient:
    """Unified client to access LLMs through multiple providers."""

    def __init__(self, config: ClientConfig):
        self.config = config

        if config.primary not in config.providers:
            raise ConfigError(f"Primary '{config.primary}' does not exist in providers")
        for fb in config.fallback:
            if fb not in config.providers:
                raise ConfigError(f"Fallback '{fb}' does not exist in providers")

        self.primary = ProviderFactory.create(config.providers[config.primary])
        self.fallbacks = [
            ProviderFactory.create(config.providers[name]) for name in config.fallback
        ]

    @classmethod
    def from_yaml(cls, path: str | Path) -> "UnifiedClient":
        with open(path) as f:
            data = yaml.safe_load(f)
        config = ClientConfig(**data)
        return cls(config)

    @classmethod
    def from_dict(cls, data: dict) -> "UnifiedClient":
        return cls(ClientConfig(**data))

    def chat(
        self,
        prompt: str,
        *,
        system: str | None = None,
        max_tokens: int = 256,
        temperature: float = 0.7,
    ) -> ChatResponse:
        """Simple version: one prompt → one response."""
        messages: list[Message] = []
        if system:
            messages.append(Message(role="system", content=system))
        messages.append(Message(role="user", content=prompt))
        return self.chat_with_messages(messages, max_tokens=max_tokens, temperature=temperature)

    def chat_with_messages(
        self,
        messages: list[Message],
        max_tokens: int = 256,
        temperature: float = 0.7,
    ) -> ChatResponse:
        """History version: full list of messages."""
        # In this capsule we only use primary. Fallback comes in capsule 04.
        return self.primary.chat(messages, max_tokens=max_tokens, temperature=temperature)

unified_ai_client/__init__.py

# unified_ai_client/__init__.py
from .client import UnifiedClient
from .models import Message, ChatResponse, ProviderConfig, ClientConfig
from .factory import ProviderFactory
from .exceptions import (
    UnifiedClientError,
    ConfigError,
    ProviderError,
    AuthError,
    RateLimitError,
    TimeoutError,
    AllProvidersFailedError,
)

__all__ = [
    "UnifiedClient",
    "Message",
    "ChatResponse",
    "ProviderConfig",
    "ClientConfig",
    "ProviderFactory",
    "UnifiedClientError",
    "ConfigError",
    "ProviderError",
    "AuthError",
    "RateLimitError",
    "TimeoutError",
    "AllProvidersFailedError",
]

Verification: use your library

Create examples/quickstart.py outside the package:

# examples/quickstart.py
"""Verifies that the library works with a real provider (OpenAI)."""
from unified_ai_client import UnifiedClient

CONFIG = {
    "primary": "openai-mini",
    "providers": {
        "openai-mini": {
            "name": "openai-mini",
            "type": "openai",
            "model": "gpt-4o-mini",
            "api_key_env": "OPENAI_API_KEY",
            "price_input_per_1m": 0.15,
            "price_output_per_1m": 0.60,
        }
    },
}

client = UnifiedClient.from_dict(CONFIG)
response = client.chat(
    "Explain REST in a single sentence.",
    max_tokens=80,
)

print(f"\nProvider: {response.provider}")
print(f"Model:    {response.model}")
print(f"Text:     {response.text}")
print(f"Tokens:   {response.tokens_input} input + {response.tokens_output} output")
print(f"Duration: {response.duration_ms}ms")
print(f"Cost:     ${response.cost_usd:.6f}" if response.cost_usd else "Cost:     n/a")

Run:

export OPENAI_API_KEY=sk-...
python examples/quickstart.py

Expected output:

Provider: openai-mini
Model:    gpt-4o-mini
Text:     REST is an architectural style for designing APIs that uses standard HTTP to create, read, update, and delete resources identified by URLs.
Tokens:   18 input + 32 output
Duration: 1843ms
Cost:     $0.000022

Verification with YAML

Create examples/clients.yaml:

primary: openai-mini
fallback: []

providers:
  openai-mini:
    name: openai-mini
    type: openai
    model: gpt-4o-mini
    api_key_env: OPENAI_API_KEY
    price_input_per_1m: 0.15
    price_output_per_1m: 0.60

And examples/quickstart_yaml.py:

from unified_ai_client import UnifiedClient
client = UnifiedClient.from_yaml("examples/clients.yaml")
print(client.chat("Just say hi").text)

OpenRouter configuration (same adapter)

providers:
  openrouter-mistral:
    name: openrouter-mistral
    type: openai_compatible
    model: mistralai/mistral-7b-instruct
    api_key_env: OPENROUTER_API_KEY
    base_url: https://openrouter.ai/api/v1
    price_input_per_1m: 0.07
    price_output_per_1m: 0.07

Change primary: openrouter-mistral and run the script again. Same code, different provider. This is the magic the entire path aims to achieve.


Local Ollama configuration

providers:
  ollama-mistral:
    name: ollama-mistral
    type: openai_compatible
    model: mistral
    base_url: http://localhost:11434/v1
    # Ollama doesn't require an api_key, but the SDK requires one
    api_key_env: OLLAMA_DUMMY_KEY
export OLLAMA_DUMMY_KEY=ollama   # any value
# Make sure Ollama is running + the model is downloaded
ollama pull mistral
ollama serve  # if it's not running as a daemon

python examples/quickstart.py  # with primary: ollama-mistral

Common traps

Trap 1 — "I forgot to activate the venv and pip install -e . installed globally." This pollutes your global Python. Always activate the venv before installing.

Trap 2 — "I changed code but when importing I see the old version." pip install -e . installs in "editable mode" — changes are reflected. If not, check that your IDE/notebook is using the correct venv.

Trap 3 — "Ollama gives me model not found." The model name must match what you have downloaded. Run ollama list to see your models and use exactly that name in the config.

Trap 4 — "I have circular import errors." It happens when models.py imports from adapters/ and vice versa. Keep models.py as a leaf of the import graph (it doesn't import from other files in the package).

Trap 5 — "My cost_usd comes out None." It's only calculated if you configured price_input_per_1m and price_output_per_1m. It's null otherwise.


Exercise

Modify the quickstart to:

  1. Pass a system prompt ("You are an assistant that responds in Hemingway's style")
  2. Print the full JSON of raw_response in addition to the text
  3. Handle the case where OPENAI_API_KEY is not set and show a useful message
See solution
# examples/quickstart_advanced.py
import os
import json
from unified_ai_client import UnifiedClient, ConfigError

if not os.environ.get("OPENAI_API_KEY"):
    print("Error: set the OPENAI_API_KEY variable before running this script.")
    exit(1)

CONFIG = {
    "primary": "openai-mini",
    "providers": {
        "openai-mini": {
            "name": "openai-mini",
            "type": "openai",
            "model": "gpt-4o-mini",
            "api_key_env": "OPENAI_API_KEY",
            "price_input_per_1m": 0.15,
            "price_output_per_1m": 0.60,
        }
    },
}

try:
    client = UnifiedClient.from_dict(CONFIG)
    response = client.chat(
        "Describe the ocean in 3 sentences.",
        system="You are Ernest Hemingway. Short sentences. Concrete images. No adverbs.",
        max_tokens=120,
    )

    print(f"Text:\n{response.text}\n")
    print(f"Raw response (JSON):")
    print(json.dumps(response.raw_response, indent=2, default=str))
except ConfigError as e:
    print(f"Config error: {e}")

Summary

You have:

  • ✅ An installable Python package (pip install -e .)
  • ✅ Pydantic models for inputs/outputs
  • ✅ Typed exceptions for error handling
  • ✅ An abstract BaseAdapter + OpenAIAdapter that covers OpenAI, OpenRouter, Ollama, LM Studio
  • ✅ A ProviderFactory that creates adapters from config
  • ✅ A UnifiedClient with chat() and chat_with_messages()
  • ✅ Loading from dict and from YAML
  • ✅ A quickstart working with a real provider

Checkpoint: if python examples/quickstart.py returns an OpenAI response with correct tokens and cost, you're ready.


Next capsule

04 — Fallback strategy. So far your client uses only primary. We're going to add automatic fallback with differentiated error handling (rate limit → retry, auth error → no retry, network error → next provider). We also see a simple circuit breaker to avoid hammering a downed provider.


Resources

  1. setuptools — pyproject.toml — modern Python package config.
  2. pip install -e (editable mode) — for library development.
  3. OpenAI Python SDK — types — the errors we catch.
  4. Pydantic v2 — Validators — to add custom validation to your models.