src/id/config.py

"""Configuration loading: pydantic models + tier/provider resolution.

Reads ``config.toml`` (providers, tiers, engine knobs, profiles) and validates
it on load. Secrets come from environment variables named by ``api_key_env``.
"""

from __future__ import annotations

import os
from pathlib import Path

try:  # Python 3.11+
    import tomllib
except ModuleNotFoundError:  # Python 3.10 (e.g. default HF Spaces image)
    import tomli as tomllib  # type: ignore[no-redef]

from pydantic import BaseModel, Field


class ProviderConfig(BaseModel):
    base_url: str
    api_key_env: str
    default_headers: dict[str, str] = Field(default_factory=dict)

    def api_key(self) -> str:
        # Many OpenAI-compatible endpoints accept any non-empty key. Fall back
        # to a placeholder so local/dummy endpoints work without an env var.
        return os.environ.get(self.api_key_env, "") or "sk-no-key-required"


class TierConfig(BaseModel):
    provider: str
    model: str
    temperature: float = 0.7
    top_p: float | None = None
    max_tokens: int | None = None


class EngineConfig(BaseModel):
    regenerate_retries: int = 2
    best_of_n: int = 3
    solver_max_attempts: int = 2
    request_timeout: float = 120.0
    max_retries: int = 3


class Config(BaseModel):
    providers: dict[str, ProviderConfig]
    tiers: dict[str, TierConfig]
    engine: EngineConfig = Field(default_factory=EngineConfig)
    profiles: dict[str, dict[str, dict[str, object]]] = Field(default_factory=dict)

    root: Path = Field(default_factory=Path.cwd, exclude=True)

    # -- resolution helpers -------------------------------------------------

    def resolve_tier(self, tier: str) -> tuple[TierConfig, ProviderConfig]:
        if tier not in self.tiers:
            raise KeyError(f"unknown tier {tier!r}; known: {sorted(self.tiers)}")
        tcfg = self.tiers[tier]
        if tcfg.provider not in self.providers:
            raise KeyError(
                f"tier {tier!r} points at unknown provider {tcfg.provider!r}"
            )
        return tcfg, self.providers[tcfg.provider]

    def with_profile(self, profile: str | None) -> Config:
        """Return a copy with a named profile's tier overrides applied."""
        if not profile:
            return self
        if profile not in self.profiles:
            raise KeyError(
                f"unknown profile {profile!r}; known: {sorted(self.profiles)}"
            )
        merged = self.model_copy(deep=True)
        for tier_name, overrides in self.profiles[profile].items():
            base = merged.tiers.get(tier_name)
            data = base.model_dump() if base else {}
            data.update(overrides)
            merged.tiers[tier_name] = TierConfig(**data)
        merged.root = self.root
        return merged

    # -- paths --------------------------------------------------------------

    @property
    def worlds_dir(self) -> Path:
        return self.root / "worlds"

    @property
    def runtime_dir(self) -> Path:
        return self.root / "runtime"

    @property
    def prompts_dir(self) -> Path:
        return self.root / "prompts"

    @property
    def prices_path(self) -> Path:
        return self.root / "prices.toml"


def load_config(path: Path | None = None) -> Config:
    """Load and validate config.toml from ``path`` (default: cwd/config.toml)."""
    root = Path.cwd()
    cfg_path = path or (root / "config.toml")
    if not cfg_path.exists():
        raise FileNotFoundError(f"config not found: {cfg_path}")
    with cfg_path.open("rb") as fh:
        raw = tomllib.load(fh)
    cfg = Config.model_validate(raw)
    cfg.root = cfg_path.parent.resolve()
    return cfg


def load_prices(path: Path) -> dict[str, dict[str, float]]:
    """Load optional model price table: model -> {prompt, completion} per 1k."""
    if not path.exists():
        return {}
    with path.open("rb") as fh:
        raw = tomllib.load(fh)
    models = raw.get("models", {})
    out: dict[str, dict[str, float]] = {}
    for model, prices in models.items():
        out[model] = {
            "prompt": float(prices.get("prompt", 0.0)),
            "completion": float(prices.get("completion", 0.0)),
        }
    return out