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
BaseAdapterandOpenAIAdapterwith 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:
- Pass a
systemprompt ("You are an assistant that responds in Hemingway's style") - Print the full JSON of raw_response in addition to the text
- Handle the case where
OPENAI_API_KEYis 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+OpenAIAdapterthat covers OpenAI, OpenRouter, Ollama, LM Studio - ✅ A
ProviderFactorythat creates adapters from config - ✅ A
UnifiedClientwithchat()andchat_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
- setuptools — pyproject.toml — modern Python package config.
- pip install -e (editable mode) — for library development.
- OpenAI Python SDK — types — the errors we catch.
- Pydantic v2 — Validators — to add custom validation to your models.