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

```toml

[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

```python

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

```python

@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

```python

@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

```python

@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

```python

@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:

```python

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`:



```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](../integrations/hosted-api.md) for the full reference.

```toml

[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

```python

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