DEV Community

Cover image for The Complete Guide to yait_aichain's Model Registry
YAIT
YAIT

Posted on

The Complete Guide to yait_aichain's Model Registry

Most AI frameworks make you hardcode provider details into every function call. You write openai.ChatCompletion.create(model="gpt-4") in forty places, then OpenAI deprecates the endpoint. You spend a Tuesday afternoon doing find-and-replace across your codebase. Two months later, you want to add Anthropic. Another Tuesday afternoon, gone.

The yait_aichain Model Registry exists so you never lose another Tuesday.

It's a single abstraction layer that maps logical model names to provider-specific configurations. You reference a model by its registry key — "openai/gpt4o" or "anthropic/claude-sonnet" — and the registry handles resolution, parameter defaults, and provider routing. Change one line in the registry, and every Skill, Chain, and Agent that references that model updates automatically.

This guide covers everything: how the registry works internally, how to configure it, how to register custom models, and how to use environment variables to keep secrets out of your source code.


How the Registry Fits Into the Architecture

The Model class is one of five core primitives in yait_aichain:

from yait_aichain import Model, Skill, Chain, Pool, Agent
Enter fullscreen mode Exit fullscreen mode

A Skill performs a single LLM operation. A Chain sequences Skills. A Pool runs them in parallel. An Agent orchestrates dynamically. None of these do anything without a Model — it's the connection between your logic and the actual inference provider.

When you instantiate a Model, the registry resolves the model key, injects default parameters, and returns a configured object ready to pass into a Skill:

model = Model("openai/gpt4o")
skill = Skill(model=model, instruction="Summarize the input text.")
result = skill.run("Paste your long document here.")
Enter fullscreen mode Exit fullscreen mode

skill.run() accepts a string and returns a string synchronously. The Skill never knows or cares which provider is behind the model — it sends a prompt, gets a response. That separation is the entire point.


The Built-In Model Registry

Out of the box, yait_aichain ships with a pre-configured registry covering the most common providers and models. You don't need to set up anything to start using them. Just reference the key.

Registry Key Format

Every model key follows the pattern provider/model_name:

Key Provider Underlying Model
openai/gpt4o OpenAI gpt-4o
openai/gpt4o-mini OpenAI gpt-4o-mini
openai/gpt4-turbo OpenAI gpt-4-turbo
anthropic/claude-sonnet Anthropic claude-3-5-sonnet-20241022
anthropic/claude-haiku Anthropic claude-3-5-haiku-20241022
anthropic/claude-opus Anthropic claude-3-opus-20240229
google/gemini-pro Google gemini-1.5-pro
google/gemini-flash Google gemini-1.5-flash
mistral/mistral-large Mistral mistral-large-latest

The slash is not decorative. The segment before it tells the registry which provider adapter to load; the segment after it identifies the specific model variant. This convention means the registry can route to the correct API client without any additional configuration from you.

The built-in registry always uses fully dated model names (e.g., claude-3-5-sonnet-20241022) when the provider requires them. If you need to target a different version, override the entry as shown in the custom registration section below.

Default Parameters

Each registered model carries a set of sensible defaults:

model = Model("openai/gpt4o")
# Defaults applied:
#   temperature: 0.7
#   max_tokens: 4096
#   top_p: 1.0
Enter fullscreen mode Exit fullscreen mode

You override any of these at instantiation:

model = Model("openai/gpt4o", temperature=0.2, max_tokens=512)
Enter fullscreen mode Exit fullscreen mode

The override applies only to that instance. The registry's defaults remain untouched for the next caller.


Configuring Models via YAML

For production systems, model configuration shouldn't be scattered across Python files. yait_aichain supports YAML-based configuration that defines your full model setup in one place.

The YAML Schema

A valid configuration file follows this structure:

version: "2.0"

models:
  openai/gpt4o:
    provider: openai
    model_name: gpt-4o
    parameters:
      temperature: 0.7
      max_tokens: 4096
      top_p: 1.0

  anthropic/claude-sonnet:
    provider: anthropic
    model_name: claude-3-5-sonnet-20241022
    parameters:
      temperature: 0.7
      max_tokens: 4096

  custom/my-finetuned:
    provider: openai
    model_name: ft:gpt-4o:my-org:custom-model:abc123
    parameters:
      temperature: 0.3
      max_tokens: 2048
Enter fullscreen mode Exit fullscreen mode

The version field is required and must be "2.0" for the current release. The models block is a dictionary where each key becomes the registry key you'll reference in code.

Loading a YAML Configuration

from yait_aichain import Model

Model.load_registry("path/to/models.yaml")

model = Model("custom/my-finetuned")
Enter fullscreen mode Exit fullscreen mode

When you call load_registry(), the models defined in YAML merge with the built-in registry. If a key in your YAML matches a built-in key, your configuration wins. This lets you override defaults without forking the library.

Schema Validation

The YAML loader validates your file against the expected schema on load. If you misspell a field or provide an invalid type, you get an error immediately — not five minutes into a pipeline run when the model finally gets called.

Common validation errors:

  • Missing version field: Every config file must declare version: "2.0".
  • Invalid provider value: Must match a supported provider adapter (openai, anthropic, google, mistral, or a custom-registered provider).
  • Parameter type mismatch: temperature must be a float between 0.0 and 2.0. max_tokens must be a positive integer.

Environment Variables

API keys and provider-specific settings should never appear in YAML files or Python source code. yait_aichain reads them from environment variables with a consistent naming convention.

Tip: Set AICHAIN_LOG_LEVEL=DEBUG whenever you're configuring the registry for the first time. The registry logs every resolution step, which makes misconfiguration obvious immediately rather than at runtime. See the Debugging section for sample output.

Required Variables by Provider

Provider Environment Variable Description
OpenAI OPENAI_API_KEY Your OpenAI API key
Anthropic ANTHROPIC_API_KEY Your Anthropic API key
Google GOOGLE_API_KEY Your Google AI API key
Mistral MISTRAL_API_KEY Your Mistral API key

Optional Configuration Variables

Beyond API keys, you can control library behavior through additional environment variables:

Variable Default Description
AICHAIN_DEFAULT_MODEL openai/gpt4o The model used when no model key is specified
AICHAIN_LOG_LEVEL WARNING Logging verbosity: DEBUG, INFO, WARNING, ERROR
AICHAIN_TIMEOUT 30 Request timeout in seconds
AICHAIN_MAX_RETRIES 3 Number of retry attempts on transient failures
AICHAIN_REGISTRY_PATH None Path to a YAML config file, loaded automatically on import

AICHAIN_REGISTRY_PATH is particularly useful in containerized deployments. Instead of calling Model.load_registry() in your application code, you set the environment variable and the registry configures itself when yait_aichain is first imported:

export AICHAIN_REGISTRY_PATH=/etc/aichain/models.yaml
export OPENAI_API_KEY=sk-...
export ANTHROPIC_API_KEY=sk-ant-...
Enter fullscreen mode Exit fullscreen mode
from yait_aichain import Model, Skill

# Registry already loaded from /etc/aichain/models.yaml
model = Model("custom/my-finetuned")
skill = Skill(model=model, instruction="Extract key entities from the text.")
result = skill.run("Apple introduced the M4 chip at its May 2024 iPad Pro event.")
Enter fullscreen mode Exit fullscreen mode

No load_registry() call. No API key in sight. The environment handles it.

Using .env Files in Development

For local development, create a .env file in your project root:

OPENAI_API_KEY=sk-dev-abc123
ANTHROPIC_API_KEY=sk-ant-dev-xyz789
AICHAIN_LOG_LEVEL=DEBUG
AICHAIN_DEFAULT_MODEL=openai/gpt4o-mini
Enter fullscreen mode Exit fullscreen mode

yait_aichain automatically detects and loads .env files from the current working directory. Add .env to your .gitignore. This is non-negotiable.


Registering Custom Models Programmatically

YAML works well for static configurations. But sometimes you need to register models at runtime — maybe you're dynamically selecting a fine-tuned model based on user input, or you're integrating a self-hosted model behind a custom API.

Basic Custom Registration

from yait_aichain import Model

Model.register(
    key="custom/llama-local",
    provider="openai",  # compatible API format
    model_name="meta-llama/Llama-3-70b",
    base_url="http://localhost:8080/v1",
    parameters={
        "temperature": 0.5,
        "max_tokens": 2048,
    }
)

model = Model("custom/llama-local")
Enter fullscreen mode Exit fullscreen mode

The base_url parameter is critical for self-hosted models. If you're running vLLM, Ollama, or any OpenAI-compatible server, set provider to "openai" and point base_url to your server. The OpenAI adapter handles the rest.

Overriding Built-In Models

You can re-register a built-in key to change its behavior globally:

Model.register(
    key="openai/gpt4o",
    provider="openai",
    model_name="gpt-4o-2024-08-06",  # pin to specific version
    parameters={
        "temperature": 0.3,  # lower default for your use case
        "max_tokens": 8192,
    }
)
Enter fullscreen mode Exit fullscreen mode

After this call, every Model("openai/gpt4o") instantiation in your application uses the pinned version with your custom defaults. This is how you roll out model version updates safely: change the registry, not the call sites.


Practical Patterns

Pattern 1: Environment-Based Model Selection

Run different models in development and production without changing code:

# Development
export AICHAIN_DEFAULT_MODEL=openai/gpt4o-mini

# Production
export AICHAIN_DEFAULT_MODEL=openai/gpt4o
Enter fullscreen mode Exit fullscreen mode
from yait_aichain import Model, Skill

# Uses whatever AICHAIN_DEFAULT_MODEL points to
model = Model()
skill = Skill(model=model, instruction="Classify the support ticket by urgency.")
result = skill.run("My account has been locked for three days and I can't log in.")
Enter fullscreen mode Exit fullscreen mode

In development you're spending fractions of a cent per call with gpt-4o-mini. In production you get the full capability of gpt-4o. Same code, different bill.

Pattern 2: Multi-Provider Fallback Chain

Use the registry to define a primary and fallback model, then build a Chain that degrades gracefully when the primary provider is unavailable:

from yait_aichain import Model, Skill, Chain

primary_model = Model("openai/gpt4o")
fallback_model = Model("anthropic/claude-sonnet")

primary_skill = Skill(
    model=primary_model,
    instruction="Generate a detailed product description."
)

fallback_skill = Skill(
    model=fallback_model,
    instruction="Generate a detailed product description."
)

chain = Chain(skills=[primary_skill, fallback_skill], mode="fallback")

try:
    result = chain.run("Noise-cancelling wireless headphones, over-ear, 30hr battery.")
except Exception as e:
    print(f"Both providers failed: {e}")
Enter fullscreen mode Exit fullscreen mode

In "fallback" mode, the Chain runs primary_skill first. If it raises a provider error, the Chain automatically retries with fallback_skill. Because both models are resolved through the registry, switching providers later means changing one key — not rewriting API calls.

Pattern 3: Centralized YAML Config for Team Projects

In a shared repository, put your model configuration in a checked-in YAML file (without secrets) and let environment variables supply the keys:

# config/models.yaml
version: "2.0"

models:
  project/summarizer:
    provider: openai
    model_name: gpt-4o
    parameters:
      temperature: 0.3
      max_tokens: 1024

  project/classifier:
    provider: anthropic
    model_name: claude-3-5-haiku-20241022
    parameters:
      temperature: 0.0
      max_tokens: 128

  project/generator:
    provider: google
    model_name: gemini-1.5-pro
    parameters:
      temperature: 0.8
      max_tokens: 4096
Enter fullscreen mode Exit fullscreen mode
from yait_aichain import Model, Skill

Model.load_registry("config/models.yaml")

summarizer = Skill(
    model=Model("project/summarizer"),
    instruction="Summarize in 3 bullet points."
)

classifier = Skill(
    model=Model("project/classifier"),
    instruction="Classify as positive, negative, or neutral."
)

result = summarizer.run("Your input text here.")
Enter fullscreen mode Exit fullscreen mode

Every team member uses the same model configurations. Nobody accidentally runs gpt-3.5-turbo when the project requires gpt-4o. The YAML file is the single source of truth for model settings; the code is the single source of truth for behavior.


Debugging Registry Issues

When something goes wrong, set AICHAIN_LOG_LEVEL=DEBUG and check the output. The registry logs every resolution step:

DEBUG:aichain.registry: Resolving model key 'openai/gpt4o'
DEBUG:aichain.registry: Found in YAML override config
DEBUG:aichain.registry: Applied parameters: temperature=0.3, max_tokens=8192
DEBUG:aichain.registry: Provider adapter: openai
DEBUG:aichain.registry: Model instance created successfully
Enter fullscreen mode Exit fullscreen mode

Common issues and their fixes:

  • "Model key not found" — You're referencing a key that doesn't exist in the built-in registry or your custom config. Check for typos. Keys are case-sensitive.
  • "Provider API key not set" — The environment variable for the provider is missing. Set OPENAI_API_KEY, ANTHROPIC_API_KEY, etc.
  • "Connection refused on base_url" — For self-hosted models, verify the server is running and the URL is correct. Include the /v1 suffix if the server expects it.
  • "YAML validation error" — Run your YAML through a linter first. The most common culprit is tabs instead of spaces, or misaligned nested keys.

What Changed in Version 2.0

If you're migrating from an earlier version, the registry saw significant changes in 2.0:

  • The version: "2.0" field is now required in YAML config files. Files without it will fail validation.
  • Model keys now use the provider/model format consistently. Old-style keys without the slash prefix are no longer supported — a 1.x key like "gpt4o" must become "openai/gpt4o".
  • Model.load_registry() merges instead of replacing. In 1.x, loading a YAML file wiped the built-in registry. In 2.0, your definitions layer on top.
  • .env file auto-loading was added. You no longer need a third-party library to load them.
  • AICHAIN_REGISTRY_PATH is new. It enables zero-code registry configuration for containerized deployments.

A typical 1.x instantiation looked like this:

# yait_aichain 1.x — no longer supported
model = Model("gpt4o")
Enter fullscreen mode Exit fullscreen mode

The 2.0 equivalent:

# yait_aichain 2.0
model = Model("openai/gpt4o")
Enter fullscreen mode Exit fullscreen mode

Search your codebase for Model(" and confirm every key contains a slash. That's the entire migration for most projects.


Summary

The Model Registry is the configuration backbone of yait_aichain. A consistent Model("provider/name") interface, centralized YAML configuration, environment-driven secrets, and clean override semantics mean that when a provider deprecates a model version, you change one registry entry and move on. Thirty seconds, not a Tuesday afternoon.

Top comments (0)