AIX: Artificial Intelligence eXtensions

A clean, pythonic facade for common AI operations that abstracts away provider-specific details and complexities.

Philosophy: Make AI interactions as simple and intuitive as Python itself, while maintaining the power and flexibility needed for production use.

Installation

pip install aix

aix is a headless facade and keeps its hard dependencies minimal — the base install only pulls in litellm (the provider backend) and config2py (for reading API keys / config). Optional capabilities live in extras and are imported lazily, raising a clear install hint if their backend is missing:

pip install "aix[image]"      # image generation helpers (Pillow)
pip install "aix[video]"      # local text-to-video (diffusers, torch)
pip install "aix[discovery]"  # model discovery via OpenRouter / Ollama (requests)
pip install "aix[google]"     # legacy Google Gemini backend
pip install "aix[openai]"     # legacy OpenAI (oa) backend
pip install "aix[all]"        # all optional backends

Quick Start

from aix import chat, embeddings, prompt_func, models
from aix import generate_image, text_to_speech, transcribe

# Simple chat
response = chat("What is 2+2?")
print(response)  # "The answer is 4."

# Create prompt-based functions
translate = prompt_func("Translate to French: {text}")
result = translate(text="Hello world")
print(result)  # "Bonjour le monde"

# Get embeddings
vecs = list(embeddings(["hello", "world"]))
print(len(vecs))  # 2

# Generate images
image = generate_image("A serene mountain landscape")
image.save("landscape.png")

# Text to speech
audio = text_to_speech("Hello, world!")
audio.save("hello.mp3")

# Speech to text
text = transcribe("recording.mp3")

# Discover available models
models.discover()
print(list(models)[:5])  # ['openai/gpt-4o', 'openai/gpt-4o-mini', ...]

Configuration

aix works with zero configuration — every function ships with sensible default models. When you want to change which model a function uses by default (and related parameters like temperature, image size, or TTS voice), there is a single source of truth: aix.config.

Defaults are resolved in layers, highest precedence first:

1. Explicit call argument — chat("hi", model="...") always wins.
2. Runtime override — aix.configure(...) (persistent) or aix.using(...) (scoped).
3. Environment variables — AIX_CHAT_MODEL, AIX_CHAT_TEMPERATURE, AIX_EMBEDDING_MODEL, AIX_IMAGE_MODEL, AIX_TTS_MODEL, AIX_TTS_VOICE, AIX_TRANSCRIPTION_MODEL, …
4. User config file — a TOML file (see below).
5. Shipped defaults — the built-in values.

import aix

# Inspect the active configuration
cfg = aix.get_config()
cfg.chat.model  # 'gpt-4.1-mini'
cfg.embeddings.model  # 'text-embedding-3-small'

# Persistently change a default (everything below an explicit arg still respects it)
aix.configure(chat_model="anthropic/claude-sonnet-4", chat_temperature=0.2)

# Scoped override that restores on exit
with aix.using(chat_model="openai/gpt-4o-mini"):
    aix.chat("quick + cheap question")
# back to the configured default here

Semantic aliases

Pass an intent-level name instead of a concrete model id. Shipped aliases: fast (cheap/low-latency), best (highest quality), cheap (cheapest).

aix.chat("hard question", model="best")  # -> resolves via the alias table
aix.chat("quick question", model="fast")

# Add or override aliases (merges with the shipped set)
aix.configure(aliases={"smart": "anthropic/claude-sonnet-4"})
aix.resolve_model("smart")  # 'anthropic/claude-sonnet-4'

Aliases share a namespace with literal model ids: a name that is not a registered alias (e.g. "gpt-4o") is passed through unchanged. Inspect available aliases via aix.get_config().aliases.

Config file

Create a TOML file at your platform's app-config dir (or point AIX_CONFIG_FILE at any path):

[chat]
model = "gpt-4.1-mini"
temperature = 1.0

[embeddings]
model = "text-embedding-3-small"

[image]
model = "dall-e-3"
size = "1024x1024"

[audio]
tts_model = "gpt-4o-mini-tts"
tts_voice = "alloy"
transcription_model = "whisper-1"

[aliases]              # semantic names -> concrete ids (override the shipped set)
fast = "gpt-4.1-mini"
best = "anthropic/claude-sonnet-4"

Find the active path with aix.config.config_file_path(). Environment variables override file values; explicit call arguments override everything.

Providing credentials

Every key-requiring function (chat, embeddings, generate_image, text_to_speech, transcribe, generate_video, ...) resolves its API key through one documented path. The resolution order, highest precedence first:

1. Explicit argument — chat("hi", api_key="sk-...") (keyword-only, on every key-requiring function).
2. Provider environment variable — the canonical name for the model's provider (see the table below). A project .env is softly discovered when python-dotenv is installed; an explicitly exported variable always wins over a .env value.
3. AIX config store — a per-user store in your app-config folder, keyed by the canonical env-var name (so the store and the env layer share one namespace).
4. (Interactive only) prompt + persist — in a REPL, AIX can ask once and save the key to the config store.

The provider is inferred from the model id, so you rarely name it yourself:

import aix

aix.chat("Hello", model="gpt-4o")  # uses OPENAI_API_KEY
aix.chat("Hi", model="claude-sonnet-4")  # uses ANTHROPIC_API_KEY
aix.chat("Yo", api_key="sk-...")  # explicit key wins

# Check what's discoverable (values are never shown — availability only):
aix.check_keys()
# {'openai': {'available': True, 'env_vars': ['OPENAI_API_KEY'], 'source': 'env'}, ...}

When a required key is missing, you get an actionable MissingCredentialError naming which key, how to set it, and where to get one — instead of a cryptic provider error:

>>> aix.chat("hi", model="gpt-4o")
aix.credentials.MissingCredentialError: No API key found for provider 'openai'
(model 'gpt-4o'). Set OPENAI_API_KEY via one of:
  - export it:  export OPENAI_API_KEY=...
  - add it to a .env file in your project, or
  - store it in the AIX config store under the key 'OPENAI_API_KEY'.
Get a key at: https://platform.openai.com/api-keys

Provider → environment variable

The full mapping is aix.PROVIDER_ENV_VARS. Keys are never logged, and the config store lives in your user app-config folder — never in the repo.

Core Features

1. Chat Interface

Clean, provider-agnostic chat completions:

from aix import chat

# Simple text prompt
response = chat("Explain quantum computing in one sentence")

# With specific model
response = chat("Hello!", model="gpt-4o-mini")

# With message history
messages = [
    {"role": "user", "content": "My name is Alice"},
    {"role": "assistant", "content": "Nice to meet you, Alice!"},
    {"role": "user", "content": "What's my name?"},
]
response = chat(messages, model="gpt-4o")

# Streaming responses
for chunk in chat("Count to 5", stream=True):
    print(chunk, end="", flush=True)

# Stateful conversations
from aix import chat_with_history

session = chat_with_history("You are a helpful math tutor")
response = session.send("What is 2+2?")
response = session.send("And if I add 3 to that?")  # Remembers context

chat() returns only the first choice's text. When you need the rest of the response (every one of n choices, token usage, why each choice stopped, the model that actually served the call), use complete(), which takes the same arguments and returns a Completion:

from aix import complete, Usage

c = complete("Name a color", n=3, temperature=1.0)
c.texts  # ('Blue', 'Green', 'Red')     -- all n choices
c.finish_reasons  # ('stop', 'stop', 'stop')     -- 'length' means truncated by max_tokens
c.usage  # Usage(prompt_tokens=11, completion_tokens=6, total_tokens=17, cached_prompt_tokens=0, ...)
c.model  # the model id the provider reports having served
c.text  # first choice, i.e. what chat() would return
c.raw  # the untouched LiteLLM response

# Usage values add up, for cost reporting. usage is None when a provider reports none.
cs = [complete(q) for q in questions]
total = sum((c.usage for c in cs if c.usage), Usage())

To ask the same question several times (agreement, stability, majority vote), use sample(). It sends one request with provider-side n when LiteLLM reports that the model supports it, and otherwise makes n independent requests concurrently. It records which path it took, because the two are not the same statistical object:

import json
from aix import sample

s = sample(
    'Is a tomato a fruit? Reply as JSON {"answer": bool}',
    n=5,
    temperature=1.0,
    parse=json.loads,
)
s.strategy  # 'provider_n', 'independent' or 'mixed'
s.payloads  # parsed answers (None where parsing failed)
[x.finish_reason for x in s]  # 'content_filter' means refused, 'length' means truncated
s.usage  # total tokens across every request that reported usage
s.n_requests  # how many requests were made

A failed request does not discard the work already paid for. By default sample() raises SamplingError once the launched requests finish, and the partial batch is on its .samples. With on_error="record", the failure is stored on that sample's .error instead. A seed is sent as-is with a provider-side n request; each independent request gets a seed derived from (seed, index).

constrained_answer(..., n=k) now runs its k calls concurrently too (concurrency=, default from config chat.sample_concurrency).

2. Embeddings

Generate vector embeddings for semantic search and similarity:

from aix import embeddings, embed, cosine_similarity, find_most_similar

# Batch embeddings
texts = ["cat", "dog", "bird"]
vecs = list(embeddings(texts))

# Single text
vec = embed("Hello, world!")

# Compute similarity
v1 = embed("cat")
v2 = embed("kitten")
similarity = cosine_similarity(v1, v2)  # High similarity

# Find most similar documents
query = "What is machine learning?"
docs = [
    "Machine learning is a type of AI",
    "Python is a programming language",
    "Neural networks are used in deep learning",
]
results = find_most_similar(query, docs, top_k=2)
# Returns: [('Machine learning is a type of AI', 0.95), ...]

# Caching for efficiency
from aix import EmbeddingCache

cache = EmbeddingCache()
v1 = cache.embed("hello")  # API call
v2 = cache.embed("hello")  # From cache

3. Prompt-Based Functions

Transform natural language prompts into reusable Python functions:

from aix import prompt_func, prompt_to_text, prompt_to_json

# Simple text generation
summarize = prompt_func("Summarize this text: {text}")
summary = summarize(text="Long article...")

# Structured output
extract_person = prompt_func(
    "Extract person information from: {text}",
    output_schema={"name": str, "age": int, "email": str},
)
result = extract_person(text="Contact John at john@example.com. He is 30 years old.")
# Returns: {'name': 'John', 'age': 30, 'email': 'john@example.com'}

# Multiple parameters
compare = prompt_func("Compare {item1} and {item2} in terms of {aspect}")
result = compare(item1="Python", item2="JavaScript", aspect="learning curve")

# Pre-built common functions
from aix import common_funcs

summary = common_funcs.summarize(text="Long article...")
keywords = common_funcs.extract_keywords(text="Article about AI and ML")
sentiment = common_funcs.sentiment(text="I love this product!")

# Create custom collections
from aix import PromptFuncs

my_funcs = PromptFuncs(model="gpt-4o")
my_funcs.add("analyze", "Analyze this code: {code}")
my_funcs.add("fix_bugs", "Fix bugs in: {code}")

result = my_funcs.analyze(code="def foo(): return bar")

Enforced, validated structured output

output_schema above is flat and only requested in the prompt; nothing checks the reply against it. For anything nested (objects inside objects, arrays of objects, enums, required vs optional), pass response_schema=, either a Pydantic model or a JSON Schema dict. It works on complete(), sample() and prompt_func():

from typing import Literal
from pydantic import BaseModel
from aix import complete, prompt_func


class Span(BaseModel):
    start: int
    end: int


class Judgement(BaseModel):
    value: Literal["yes", "no", "unsure"]
    confidence: float
    evidence: list[Span]


judge = prompt_func("Does this text claim {claim}? {text}", response_schema=Judgement)
judge(
    claim="...", text="..."
)  # -> Judgement(value='yes', confidence=0.8, evidence=[...])

c = complete("...", response_schema=Judgement)
c.value  # the validated Judgement (raises SchemaValidationError if still invalid)
c.schema_mode  # 'strict' | 'schema' | 'json_mode' -- how the provider enforced it
c.repairs  # extra requests made to fix invalid replies (their tokens are in c.usage)

The provider enforces the schema as strongly as the model allows:

• strict: a strict json_schema response format, which constrains decoding. Used when LiteLLM's supports_response_schema(model) is true and the schema is strict-compatible (every object closes additionalProperties and requires all of its properties; Pydantic models are converted to that form automatically).
• schema: the same format, non-strict, for a schema that is not strict-compatible.
• json_mode: JSON mode plus the schema written into the prompt, for models without schema support (or whose support is unknown).
• prompt: the schema in the prompt only, for models LiteLLM knows accept no response_format at all.

A schema whose root is not an object (a list, a union, a RootModel) is sent wrapped as {"value": ...}, since providers only constrain objects. It is unwrapped for you after validation.

Every reply is also validated client-side. An invalid reply is re-asked up to max_repairs times (default from config chat.schema_max_repairs, which is 1), and each repair request shows the model its error. A refusal is recorded and never re-asked. A failed repair request is recorded on that choice; the other choices are kept.

4. Model Discovery & Selection

Discover and filter models across providers:

from aix import models

# Discover available models
models.discover("openrouter")  # Fetch 400+ models from OpenRouter

# List all models
all_models = list(models)

# Get specific model info
info = models["openai/gpt-4o"]
print(info.provider)  # 'openai'
print(info.context_size)  # 128000

# Filter models
openai_models = models.filter(provider="openai")
cheap_models = models.filter(
    custom_filter=lambda m: m.cost_per_token.get("input", 0) < 0.001
)
local_models = models.filter(is_local=True)

# Search models
results = models.search("gpt-4")
results = models.search("claude")

# Get recommendations
recommended = models.recommend(
    task="chat", max_cost_per_mtok=5.0, min_context_size=16000
)

# Use with chat
model = models["gpt-4o-mini"]
response = chat("Hello", model=model.id)

5. Batch Operations

Process multiple requests efficiently:

from aix import batch_chat, batch_embeddings, BatchProcessor

# Batch chat
prompts = ["What is 2+2?", "What is 3+3?", "What is 5+5?"]
results = list(batch_chat(prompts, batch_size=10, max_workers=5))

# Batch embeddings
texts = ["hello", "world", "foo", "bar"] * 100
vectors = list(batch_embeddings(texts, batch_size=20, show_progress=True))

# Generic batch processing
from aix import batch_process


def analyze(text):
    return chat(f"Analyze sentiment: {text}")


texts = ["I love it!", "It's okay", "Terrible"]
results = list(batch_process(texts, analyze, batch_size=5, retry_attempts=3))

# Stateful batch processor
processor = BatchProcessor(show_progress=True)
results = processor.process_chats(prompts)
processor.save_results("output.json")

6. Image Generation

Generate images from text descriptions:

from aix import generate_image, generate_images

# Simple image generation
image = generate_image("A serene mountain landscape at sunset")
image.save("landscape.png")

# High quality with DALL-E 3
image = generate_image(
    "Abstract art with vibrant colors", model="dall-e-3", quality="hd", style="vivid"
)

# Generate multiple variations
images = generate_images("A cute robot waving hello", n=3, size="512x512")
for i, img in enumerate(images):
    img.save(f"robot_{i}.png")

# Edit existing images
from aix import edit_image

edited = edit_image("photo.jpg", "Add a rainbow in the sky", mask_path="sky_mask.png")

# Create variations
from aix import create_variation

variations = create_variation("original.png", n=3)

7. Audio Operations

Text-to-speech and speech-to-text:

from aix import text_to_speech, transcribe, transcribe_with_timestamps

# Text to speech
audio = text_to_speech("Hello, world!")
audio.save("hello.mp3")

# Different voices
audio = text_to_speech("This is a test", voice="nova", speed=1.2)

# Transcribe audio
text = transcribe("recording.mp3")
print(text)  # "This is the transcribed text"

# With language hint
text = transcribe("spanish_audio.mp3", language="es")

# Detailed transcription with timestamps
result = transcribe_with_timestamps("lecture.mp3")
for segment in result.segments:
    print(f"[{segment['start']:.2f}] {segment['text']}")

# Translate audio to English
from aix import translate_audio

english_text = translate_audio("spanish_audio.mp3")

8. Video Generation (Coming Soon)

Video generation with provider-specific implementations:

from aix import generate_video, get_video_providers

# Check available providers
providers = get_video_providers()
print(providers)  # ['runway', 'pika', ...]

# Generate video (requires provider setup)
video = generate_video(
    "A cat walking through a garden", duration=5, resolution="1920x1080"
)
video.save("cat_video.mp4")

# Animate static image
from aix import animate_image_to_video

video = animate_image_to_video(
    "landscape.jpg", prompt="Gentle camera pan across the scene"
)

Note: Video generation requires additional provider setup (Runway, Pika, etc.) and API keys.

OpenRouter Integration

OpenRouter provides a single API key for 400+ models from 60+ providers. This is recommended for:

• Getting started quickly - One key vs. managing many
• Experimenting - Easy access to models from OpenAI, Anthropic, Google, etc.
• Production flexibility - Switch providers without code changes

Setup

1. Get API key from https://openrouter.ai
2. Set environment variable:

   export OPENROUTER_API_KEY=your-key-here
   

3. Use OpenRouter models:

   # Prefix models with 'openrouter/'
   chat("Hello", model="openrouter/openai/gpt-4o")
   chat("Hello", model="openrouter/anthropic/claude-3.5-sonnet")
   
   # Discover available models
   models.discover("openrouter")
   

All standard AIX features work with OpenRouter models.

Architecture

AIX follows the i2mint philosophy of clean, functional interfaces:

• Mapping interfaces for collections (models, registries)
• Functional approach over verbose OOP
• Lazy evaluation via generators where appropriate
• Protocol-based design for flexibility

Backend: LiteLLM

AIX uses LiteLLM as the backend for provider interactions, but wraps it in clean, pythonic interfaces. Users never need to interact with LiteLLM directly.

Supported Providers (via LiteLLM):

• OpenAI
• Anthropic (Claude)
• Google (Gemini)
• Mistral
• Cohere
• And 100+ more

import aix does not touch the network

LiteLLM fetches its model-cost map over HTTPS as a side effect of being imported. aix therefore resolves LiteLLM on the first call rather than at import time, so importing aix (or any of its submodules) opens no outbound connection: a server cold start doesn't block on it, an offline or restricted-egress deployment doesn't eat a connect timeout, and a test suite that merely imports aix stays hermetic. The first actual provider call pays the import, and gets the live cost map — so provider resolution is unaffected.

If you also want the first call to stay offline, set LiteLLM's own knob in your environment. aix deliberately never sets it for you — it is a whole-process setting, so the choice is the operator's:

export LITELLM_LOCAL_MODEL_COST_MAP=1

The trade-off: LiteLLM then reads the cost map bundled in its wheel, which is a snapshot. Bare model names newer than your installed litellm release may stop resolving — pass those provider-prefixed ("openai/gpt-5.2" rather than "gpt-5.2"), which works either way.

Design Patterns

From oa (OpenAI facade)

AIX builds on patterns from the oa package:OpenRouter

• Simple chat interface with smart defaults
• Template-based function creation
• Structured output support
• Batch processing capabilities

From i2 (Signature manipulation)

• Clean function signatures
• Flexible parameter handling
• Decorator-based composition

From dol (Storage abstraction)

• Mapping-based interfaces
• Persistent storage options
• Cache management

Advanced Usage

Custom Model Sources

from aix.ai_models import ModelManager, OpenRouterSource

manager = ModelManager()
source = OpenRouterSource()
models = manager.discover_from_source("openrouter")

Connector-Specific Metadata

# Get provider-specific parameters
metadata = models.get_connector_metadata("openai/gpt-4o", "openai")
# Use with native SDK: openai.ChatCompletion.create(**metadata, messages=[...])

Error Handling

from aix import chat

try:
    response = chat("Hello", model="nonexistent-model")
except Exception as e:
    print(f"Error: {e}")

Backward Compatibility

The legacy chat_funcs and chat_models interfaces are still available:

from aix import chat_funcs, chat_models

# Old style (still works)
list(chat_funcs)  # ['gpt-4o', 'gpt-4o-mini', ...]
response = chat_funcs.gpt_4o("Hello")

# Model metadata
info = chat_models["gpt-4o"]
# {'price_per_million_tokens': 5.0, 'provider': 'openai', ...}

However, the new interfaces (chat, embeddings, prompt_func, models) are recommended for new code.

Examples

Semantic Search

from aix import embed, cosine_similarity

# Build document index
docs = ["AI is the future", "Python is great", "Machine learning works"]
doc_vecs = [embed(doc) for doc in docs]

# Search
query_vec = embed("artificial intelligence")
similarities = [cosine_similarity(query_vec, dv) for dv in doc_vecs]
best_match = docs[similarities.index(max(similarities))]
print(best_match)  # "AI is the future"

Data Extraction Pipeline

from aix import prompt_func, batch_process

# Define extraction function
extract = prompt_func(
    "Extract product info from: {text}",
    output_schema={"name": str, "price": float, "category": str},
)

# Process many product descriptions
descriptions = [...]  # Your data
results = list(
    batch_process(
        descriptions, lambda d: extract(text=d), batch_size=10, show_progress=True
    )
)

Multi-Model Comparison

from aix import chat

prompt = "Explain quantum computing in one sentence"

# Try different models
for model_id in ["gpt-4o-mini", "claude-sonnet-4", "gemini-1.5-flash"]:
    response = chat(prompt, model=model_id)
    print(f"{model_id}: {response}")

Documentation

• API Reference: See docstrings in each module
• Examples: Check examples/ directory
• GitHub: https://github.com/thorwhalen/aix

Contributing

Contributions are welcome! Please see CONTRIBUTING.md for guidelines.

License

Apache 2.0

Credits

Built with:

• LiteLLM - Multi-provider backend
• i2 - Signature manipulation utilities
• dol - Storage abstraction patterns
• oa - OpenAI facade inspiration

What AIX is NOT

AIX is NOT an AI agent framework. For that, see the separate aw package. AIX focuses on foundational AI operations (chat, embeddings, etc.) that can be used to build agents and other applications.