docs/design/CLI_DESIGN.md

CLI & Provider Configuration Design

Design document for the ace CLI, model configuration system, and provider registry.

---

Implementation Status

Component	Status	Location
ModelConfig / ACEModelConfig dataclasses	Done	ace/providers/config.py
TOML persistence (ace.toml)	Done	ace/providers/config.py
.env persistence (secrets)	Done	ace/providers/config.py
Provider registry (LiteLLM delegation)	Done	ace/providers/registry.py
Model search / discovery	Done	ace/providers/registry.py
Connection validation	Done	ace/providers/registry.py
ace setup interactive wizard	Done	ace/cli/setup.py
ace models search command	Done	ace/cli/setup.py
ace validate connection test	Done	ace/cli/setup.py
Lazy imports (fast CLI startup)	Done	ace/__init__.py, ace/providers/__init__.py, ace/providers/registry.py

---

Goals

1. Zero-friction onboarding — ace setup walks the user from nothing to a working config in under a minute.
2. Secrets / config separation — ace.toml (committable, no secrets) + .env (gitignored API keys). Teams share model choices without leaking credentials.
3. Per-role model selection — different models for Agent, Reflector, and Skill Manager to optimise cost vs quality.
4. Provider agnosticism — any model string LiteLLM supports works. No provider-specific code in the config layer.
5. Fast CLI startup — heavy dependencies (LiteLLM, instructor, openai) are lazily imported. The ace command starts in ~50ms, not ~2s.

---

Architecture Overview

pyproject.toml
  [project.scripts]
  ace = "ace.cli.setup:main"        # entry point

ace/cli/
  __init__.py
  setup.py          # CLI commands + interactive wizard

ace/providers/
  __init__.py        # lazy re-exports
  config.py          # ModelConfig, ACEModelConfig, TOML + .env I/O
  registry.py        # provider detection, model search, connection validation
  pydantic_ai.py     # resolve_model, settings_from_config (PydanticAI model resolution)

The CLI layer (ace/cli/) depends on:

• ace/providers/config.py — always (lightweight, no heavy deps)
• ace/providers/registry.py — only when validating or searching (imports LiteLLM lazily)

The config layer has zero heavy dependencies — it uses only tomllib (stdlib), pathlib, and dataclasses.

---

Configuration Model

Two-file split

File	Contains	Git	Written by
ace.toml	Model names, temperature, max_tokens, extra_params	Commit	save_config()
.env	OPENAI_API_KEY=sk-..., ANTHROPIC_API_KEY=sk-ant-...	Gitignore	save_env_var()

ace.toml format

[default]
model = "gpt-4o-mini"

[agent]
model = "claude-sonnet-4-20250514"
max_tokens = 4096

[reflector]
model = "gpt-4o-mini"
temperature = 0.2

Roles without an explicit section inherit from [default]. Only non-default values are written (e.g. temperature is omitted when it equals 0.0).

Config discovery

find_config(start) walks up from start to the filesystem root looking for ace.toml. This supports monorepos where the config lives at the project root but commands run from subdirectories.

Loading in code

from ace import ACELiteLLM

# Option 1: Load from ace.toml + .env (created by `ace setup`)
ace = ACELiteLLM.from_setup()

# Option 2: Explicit model, keys from environment
ace = ACELiteLLM.from_model("gpt-4o-mini")

# Option 3: Full config object
from ace import ACEModelConfig, ModelConfig
ace = ACELiteLLM.from_config(ACEModelConfig(
    default=ModelConfig(model="gpt-4o-mini"),
    agent=ModelConfig(model="claude-sonnet-4-20250514"),
))

---

Data Types

ModelConfig

@dataclass
class ModelConfig:
    model: str                              # LiteLLM model string
    temperature: float = 0.0
    max_tokens: int = 2048
    extra_params: dict[str, Any] | None = None

Serialises to/from a TOML section. to_dict() omits default values to keep the file clean.

ACEModelConfig

@dataclass
class ACEModelConfig:
    default: ModelConfig                    # required — used as fallback
    agent: ModelConfig | None = None        # overrides default for Agent role
    reflector: ModelConfig | None = None    # overrides default for Reflector role
    skill_manager: ModelConfig | None = None  # overrides default for Skill Manager role

for_role(role) returns the role-specific config or falls back to default.

ValidationResult

@dataclass
class ValidationResult:
    success: bool
    model: str = ""
    provider: str = ""
    latency_ms: int = 0
    error: str = ""

Returned by validate_connection(). On success, includes the provider name and round-trip latency. On failure, includes a human-readable error string.

ModelInfo

@dataclass
class ModelInfo:
    model: str
    provider: str
    max_input_tokens: int | None = None
    max_output_tokens: int | None = None
    input_cost_per_m: float | None = None   # cost per million tokens
    output_cost_per_m: float | None = None
    key_found: bool = False                 # are the required env vars set?

Returned by search_models(). Pricing is per million tokens (converted from LiteLLM's per-token values).

---

Provider Registry

All provider logic is delegated to LiteLLM. The registry module (ace/providers/registry.py) wraps LiteLLM with a stable API:

Function	What it does	Makes API calls?
get_provider(model)	Detect provider from model string	No
get_missing_keys(model)	List required env vars that are unset	No
keys_are_set(model)	Check if all required env vars exist	No
validate_connection(model, api_key?)	Send a 3-token test call	Yes (tiny)
search_models(query, provider?, limit?)	Search LiteLLM's static model registry	No
suggest_models(typo, limit?)	Fuzzy-match model names for typo correction	No

LiteLLM lazy import

LiteLLM takes ~1.5s to import. The registry defers import until first use:

def _litellm():
    global _litellm_mod
    try:
        return _litellm_mod
    except NameError:
        pass
    import litellm as _mod
    _litellm_mod = _mod
    return _mod

All registry functions call _litellm() instead of using a top-level import.

Provider key mapping

PROVIDER_KEY_ENV maps provider names to the environment variables they require:

Provider	Required env vars
openai	OPENAI_API_KEY
anthropic	ANTHROPIC_API_KEY
azure	AZURE_API_KEY
gemini	GEMINI_API_KEY
deepseek	DEEPSEEK_API_KEY
groq	GROQ_API_KEY
bedrock	AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY + AWS_REGION_NAME
bedrock_converse	AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY + AWS_REGION_NAME
vertex_ai	GOOGLE_APPLICATION_CREDENTIALS
cohere	COHERE_API_KEY
mistral	MISTRAL_API_KEY
openrouter	OPENROUTER_API_KEY
together_ai	TOGETHERAI_API_KEY
fireworks_ai	FIREWORKS_AI_API_KEY
replicate	REPLICATE_API_KEY
huggingface	HUGGINGFACE_API_KEY
perplexity	PERPLEXITYAI_API_KEY
anyscale	ANYSCALE_API_KEY

For multi-key providers (bedrock), all listed vars must be set for _quick_key_check to report True.

This mapping is preferred over LiteLLM's validate_environment() because:

• LiteLLM's response can be inaccurate for certain providers (e.g. bedrock_converse)
• Our mapping is used for both the Key column in ace models and the key prompting in ace setup

---

CLI Commands

Entry point

Defined in pyproject.toml:

[project.scripts]
ace = "ace.cli.setup:main"

main() uses argparse with subcommands:

ace setup [--dir DIR]        Interactive configuration wizard
ace models [QUERY] [--provider P] [--limit N]   Search model catalog
ace validate MODEL           Test a model connection
ace config                   Show current configuration

ace setup

Interactive wizard flow:

1. Load existing .env (if present)
2. Check for existing ace.toml
   - If found: show current config, ask "Reconfigure?"
   - If declined: return existing config
3. Step 1: Choose your model
   - Prompt for model name
   - Detect provider via get_provider()
   - Validate connection
     - If auth fails: prompt for missing keys, retry
     - If model not found: suggest alternatives, re-prompt
     - If success: continue
4. Step 2: Role assignment
   - Ask "Use this model for all roles?" (default: yes)
   - If no: prompt for each role (Agent, Reflector, Skill Manager)
     - Enter = keep default (skip validation)
     - Different model = validate it
5. Save ace.toml and .env
6. Print configuration summary

Key behaviours:

• Validation-first: the wizard tries the connection immediately. If credentials exist in the environment (env vars, .env, AWS profiles), no prompting needed.
• Error recovery: on failed validation, env vars set during prompting are rolled back.
• Secret handling: API keys are prompted via getpass (hidden input). Non-secret values like AWS_REGION_NAME and GOOGLE_APPLICATION_CREDENTIALS use regular input() so users can see what they type.
• Typo correction: when a model is not found, suggest_models() offers alternatives.

ace models

Searches LiteLLM's static model_cost registry (no API calls):

$ ace models claude haiku

Model                                         Provider        Input $/M  Output $/M  Key
------------------------------------------------------------------------------------------
claude-haiku-4-5-20251001                     anthropic        $1.00      $5.00       ✓
us.anthropic.claude-haiku-4-5-20251001-v1:0   bedrock_converse $1.10      $5.50       ✗

• Multiple query terms are AND-matched (all must appear in the model name)
• --provider filters by LiteLLM provider name
• --limit caps results (default 20), shows total count
• Key column uses _quick_key_check() — checks env vars only, no API calls

ace validate

Sends a minimal LLM call (3 tokens: "Say 'ok'") to verify:

• API key authentication
• Model availability at the provider
• Network connectivity

$ ace validate gpt-4o-mini
✓ Connected! (gpt-4o-mini via openai, 203ms)

On failure, suggests similar model names if the model wasn't found.

All subcommands (models, validate, config) use _load_project_dotenv() which finds .env relative to ace.toml (via find_config()), not just CWD.

ace config

Displays the current configuration from ace.toml:

$ ace config
Configuration (/path/to/ace.toml)

  Role             Model
  ---------------- ---------------------------------------------
  default          gpt-4o-mini
  agent            claude-sonnet-4-20250514
  reflector        (default)
  skill_manager    (default)

Uses find_config() to locate ace.toml from the current directory upward.

ace models (no query)

When called without arguments, shows usage examples instead of dumping arbitrary results:

$ ace models
Usage: ace models <query>

Examples:
  ace models claude          All Claude models
  ace models gpt 4o          GPT-4o variants
  ace models haiku us        US-region Haiku models
  ace models --provider openai  All OpenAI models

---

Kayba CLI

The kayba entry point (ace.cli:main) provides the hosted API client. It wraps trace management, pipeline execution, insight triage, prompt generation, and integration management. See Hosted API docs for the full reference.

[project.scripts]
ace = "ace.cli.setup:main"
kayba = "ace.cli:main"
ace-mcp = "ace.integrations.mcp.server:main"

Key command groups: traces, run, insights, prompts, integrations, status, materialize, batch, setup.

---

Lazy Import Strategy

Problem

import ace eagerly imported all submodules, including litellm (1.5s). This made the CLI unusable (2s startup for a simple ace --help).

Solution

Three-layer lazy import:

1. ace/__init__.py — __getattr__-based lazy loading. TYPE_CHECKING block for IDE support, _LAZY_IMPORTS dict for runtime.

2. ace/providers/__init__.py — same pattern. Config imports are eager (lightweight), everything else is lazy.

3. ace/providers/registry.py — _litellm() helper defers import litellm until first function call.

Result: ace --help runs in ~50ms. LiteLLM is only imported when the user actually searches, validates, or sets up.

Pattern

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .heavy_module import HeavyClass  # IDE autocomplete

_LAZY_IMPORTS = {
    "HeavyClass": ("package.heavy_module", "HeavyClass"),
}

def __getattr__(name: str) -> object:
    if name in _LAZY_IMPORTS:
        module_path, attr = _LAZY_IMPORTS[name]
        import importlib
        module = importlib.import_module(module_path)
        value = getattr(module, attr)
        globals()[name] = value  # cache for subsequent access
        return value
    raise AttributeError(...)

---

File Map

ace/
  __init__.py                     # lazy re-exports for all ace symbols
  cli/
    setup.py                      # main(), run_setup(), _cmd_models(), _cmd_validate(), _cmd_config()
  providers/
    __init__.py                   # lazy re-exports (config eager, rest lazy)
    config.py                     # ModelConfig, ACEModelConfig, TOML/env I/O
    registry.py                   # provider detection, model search, validation
    pydantic_ai.py                # resolve_model, settings_from_config (PydanticAI model resolution)

Generated files:

project-root/
  ace.toml                        # model config (commit this)
  .env                            # API keys (gitignore this)

---

Known Issues

Resolved

#	Issue	Fix
1	_prompt_secret used getpass for non-secrets like AWS_REGION_NAME	Non-secret vars (AWS_REGION_NAME, GOOGLE_APPLICATION_CREDENTIALS) now use visible _prompt()
2	Error classification used fragile substring matching	Simplified: only "not found" is non-recoverable; everything else offers key prompting
3	save_env_var didn't quote values	Values now written as KEY="value"
4	_PROVIDER_KEY_ENV was private but imported externally	Renamed to PROVIDER_KEY_ENV (public)
5	v / x status symbols	Replaced with ✓ / ✗
7	ace models with no query dumped arbitrary results	Now shows usage examples instead
8	No way to inspect current config	Added ace config command
10	ace validate / ace models only loaded .env from CWD	Added _load_project_dotenv() — finds .env relative to ace.toml via find_config()
11	Unused imports asdict, field in config.py	Removed
12	Dead function _litellm_available() in registry.py	Removed
13	Unused PROVIDER_MODEL_EXAMPLES import in setup.py	Removed
14	Duplicate search_models import in setup.py	Consolidated to single top-level import

Open

#	Issue	Impact
6	No --non-interactive mode for CI/Docker	Blocks CI automation — requires a future ace setup --model MODEL --skip-validation flag
9	Per-role model selection skips validation when keeping default	Low — the default was already validated in Step 1
15	No multi-provider key setup in one pass	When reconfiguring (ace setup on an existing config), users who want different providers per role (e.g. OpenAI default + Anthropic agent + Bedrock reflector) must go through each role sequentially — keys are only prompted when a model fails validation. A future improvement should let users configure all providers and their keys upfront in a single credentials step, before role assignment begins.

---

Design Decisions

Why hand-rolled TOML serialisation?

Python's tomllib (stdlib since 3.11) only reads TOML. Writing requires tomli-w (third-party) or a manual serialiser. To avoid adding a dependency for a simple four-section config file, we hand-roll _to_toml(). The format is simple enough that this is reliable — the only complex case is extra_params (inline table).

Why our own key mapping instead of LiteLLM's?

LiteLLM's validate_environment() sometimes returns incorrect keys (e.g. suggesting AWS_BEARER_TOKEN_BEDROCK for bedrock_converse when the standard auth is AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY + AWS_REGION_NAME). Our PROVIDER_KEY_ENV mapping is simpler, auditable, and covers the common providers. For unknown providers, we fall back to LiteLLM's response or guess {PROVIDER}_API_KEY.

Why validation-first in setup?

Many users already have credentials in their environment (exported vars, AWS profiles, .env from another project). Prompting for keys before trying would waste their time. By attempting the connection first, the happy path is: type model name → instant success → done.

Why ace.toml instead of pyproject.toml [tool.ace]?

• Keeps ACE config decoupled from the Python project (ACE might be used in non-Python contexts)
• find_config() can walk up the directory tree independently
• Easier to reason about — one file, one purpose