markitect-main/markitect/llm/toml_config.py

"""
TOML-based LLM configuration: defaults, preferences, and resolution.

Config files:
  - Directory: ``<dir-with-pyproject.toml>/.markitect.toml``
  - User:      ``~/.config/markitect/config.toml``

Resolution order (highest → lowest):
  1. CLI flags (``--provider``, ``--model``)
  2. ``MARKITECT_HELPER_MODEL`` env var (model only)
  3. User preference  (``[llm.preference]`` in user config)
  4. Directory preference (``[llm.preference]`` in directory config)
  5. Directory default  (``[llm.default]`` in directory config)
  6. User default     (``[llm.default]`` in user config)
  7. Hardcoded fallback
"""

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import toml

from markitect.llm.config import find_project_root

# ── Constants ─────────────────────────────────────────────────────────────

HARDCODED_PROVIDER = "gemini"
HARDCODED_MODEL = "gemini-2.5-flash"
MODEL_ENV_VAR = "MARKITECT_HELPER_MODEL"

USER_CONFIG_DIR = Path.home() / ".config" / "markitect"
USER_CONFIG_PATH = USER_CONFIG_DIR / "config.toml"
DIR_CONFIG_NAME = ".markitect.toml"


# ── Data classes ──────────────────────────────────────────────────────────

@dataclass
class LLMLayer:
    """One layer of provider/model configuration (may be partial)."""
    provider: Optional[str] = None
    model: Optional[str] = None


@dataclass
class ResolvedLLM:
    """Fully-resolved provider + model with source attribution."""
    provider: str
    model: str
    provider_source: str
    model_source: str


# ── Read / Write / Clear ─────────────────────────────────────────────────

def _read_llm_section(path: Path, section: str) -> LLMLayer:
    """Read ``[llm.<section>]`` from a TOML file. Returns empty layer on error."""
    try:
        data = toml.load(path)
    except (OSError, toml.TomlDecodeError):
        return LLMLayer()
    llm = data.get("llm", {})
    sec = llm.get(section, {})
    return LLMLayer(
        provider=sec.get("provider"),
        model=sec.get("model"),
    )


def _write_llm_section(path: Path, section: str, layer: LLMLayer) -> None:
    """Merge ``[llm.<section>]`` into a TOML file. Creates dirs as needed."""
    path.parent.mkdir(parents=True, exist_ok=True)

    try:
        data = toml.load(path)
    except (OSError, toml.TomlDecodeError):
        data = {}

    llm = data.setdefault("llm", {})
    sec = llm.setdefault(section, {})

    if layer.provider is not None:
        sec["provider"] = layer.provider
    if layer.model is not None:
        sec["model"] = layer.model

    with open(path, "w") as f:
        toml.dump(data, f)


def _clear_llm_section(path: Path, section: str) -> bool:
    """Remove ``[llm.<section>]``. Returns True if something was cleared."""
    try:
        data = toml.load(path)
    except (OSError, toml.TomlDecodeError):
        return False

    llm = data.get("llm")
    if not isinstance(llm, dict) or section not in llm:
        return False

    del llm[section]

    # Clean up empty [llm] table.
    if not llm:
        del data["llm"]

    with open(path, "w") as f:
        toml.dump(data, f)
    return True


# ── Directory config path helper ─────────────────────────────────────────

def _dir_config_path() -> Optional[Path]:
    root = find_project_root()
    if root is None:
        return None
    return root / DIR_CONFIG_NAME


# ── Resolution ───────────────────────────────────────────────────────────

def resolve_llm(
    cli_provider: Optional[str] = None,
    cli_model: Optional[str] = None,
) -> ResolvedLLM:
    """Walk the 7-level priority chain and return a fully resolved config.

    Provider and model are resolved independently — each takes the value
    from its highest-priority source.
    """
    dir_path = _dir_config_path()

    # Build the layers (highest priority first).
    layers: list[tuple[str, LLMLayer]] = []

    # 1. CLI flags
    layers.append(("CLI flag", LLMLayer(provider=cli_provider, model=cli_model)))

    # 2. Env var (model only)
    env_model = os.environ.get(MODEL_ENV_VAR) or None
    layers.append(("env MARKITECT_HELPER_MODEL", LLMLayer(model=env_model)))

    # 3. User preference
    layers.append((
        "user preference",
        _read_llm_section(USER_CONFIG_PATH, "preference"),
    ))

    # 4. Directory preference
    if dir_path:
        layers.append((
            "directory preference",
            _read_llm_section(dir_path, "preference"),
        ))

    # 5. Directory default
    if dir_path:
        layers.append((
            "directory default",
            _read_llm_section(dir_path, "default"),
        ))

    # 6. User default
    layers.append((
        "user default",
        _read_llm_section(USER_CONFIG_PATH, "default"),
    ))

    # 7. Hardcoded
    layers.append(("hardcoded", LLMLayer(provider=HARDCODED_PROVIDER, model=HARDCODED_MODEL)))

    # Resolve provider and model independently (first non-None wins).
    provider = HARDCODED_PROVIDER
    provider_source = "hardcoded"
    model = HARDCODED_MODEL
    model_source = "hardcoded"

    for source, layer in layers:
        if layer.provider:
            provider = layer.provider
            provider_source = source
            break

    for source, layer in layers:
        if layer.model:
            model = layer.model
            model_source = source
            break

    return ResolvedLLM(
        provider=provider,
        model=model,
        provider_source=provider_source,
        model_source=model_source,
    )


def get_default_layers() -> list[tuple[str, LLMLayer]]:
    """Return only the default layers for display."""
    dir_path = _dir_config_path()
    layers: list[tuple[str, LLMLayer]] = []

    if dir_path:
        layers.append((
            f"Directory default ({DIR_CONFIG_NAME})",
            _read_llm_section(dir_path, "default"),
        ))

    layers.append((
        f"User default ({USER_CONFIG_PATH})",
        _read_llm_section(USER_CONFIG_PATH, "default"),
    ))

    layers.append((
        "Hardcoded",
        LLMLayer(provider=HARDCODED_PROVIDER, model=HARDCODED_MODEL),
    ))

    return layers


def get_preference_layers() -> list[tuple[str, LLMLayer]]:
    """Return only the preference layers for display."""
    dir_path = _dir_config_path()
    layers: list[tuple[str, LLMLayer]] = []

    layers.append((
        f"User preference ({USER_CONFIG_PATH})",
        _read_llm_section(USER_CONFIG_PATH, "preference"),
    ))

    if dir_path:
        layers.append((
            f"Directory preference ({DIR_CONFIG_NAME})",
            _read_llm_section(dir_path, "preference"),
        ))

    return layers