model-selector/model_selector/config.py

from __future__ import annotations

import json
import shutil
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable


PROVIDER_LABELS = {
    "openai": "OpenAI",
    "groq-cloud": "Groq",
    "openrouter": "OpenRouter",
}


@dataclass(frozen=True)
class ConfiguredModel:
    provider_key: str
    provider_label: str
    model_key: str
    model_id: str
    name: str
    reasoning: bool | None
    inputs: tuple[str, ...]
    context_window: int | None
    max_tokens: int | None
    alias: str | None

    @property
    def full_key(self) -> str:
        return f"{self.provider_key}/{self.model_id}"


def provider_label(provider_key: str) -> str:
    return PROVIDER_LABELS.get(provider_key, provider_key.replace("-", " ").title())


class OpenClawConfigRepository:
    def __init__(self, config_path: Path) -> None:
        self.config_path = config_path

    def load(self) -> dict[str, Any]:
        with self.config_path.open("r", encoding="utf-8") as handle:
            return json.load(handle)

    def save(self, config: dict[str, Any]) -> None:
        self.config_path.parent.mkdir(parents=True, exist_ok=True)
        if self.config_path.exists():
            backup_path = self.config_path.with_suffix(self.config_path.suffix + ".bak")
            shutil.copy2(self.config_path, backup_path)

        with tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            dir=self.config_path.parent,
            delete=False,
        ) as handle:
            json.dump(config, handle, indent=2, ensure_ascii=False)
            handle.write("\n")
            temp_path = Path(handle.name)

        temp_path.replace(self.config_path)

    def configured_models(self, config: dict[str, Any] | None = None) -> list[ConfiguredModel]:
        data = config if config is not None else self.load()
        providers = data.get("models", {}).get("providers", {})
        aliases = data.get("agents", {}).get("defaults", {}).get("models", {})
        models: list[ConfiguredModel] = []

        for provider_key, provider in providers.items():
            provider_models = provider.get("models", []) or []
            for model in provider_models:
                model_id = str(model.get("id", ""))
                full_key = f"{provider_key}/{model_id}"
                alias = aliases.get(full_key, {}).get("alias")
                models.append(
                    ConfiguredModel(
                        provider_key=provider_key,
                        provider_label=provider_label(provider_key),
                        model_key=full_key,
                        model_id=model_id,
                        name=str(model.get("name", model_id)),
                        reasoning=model.get("reasoning"),
                        inputs=tuple(model.get("input", []) or ()),
                        context_window=model.get("contextWindow"),
                        max_tokens=model.get("maxTokens"),
                        alias=alias,
                    )
                )

        return models

    def update_alias(self, full_key: str, alias: str | None) -> dict[str, Any]:
        config = self.load()
        model_aliases = self._aliases_section(config)

        if alias:
            model_aliases[full_key] = {"alias": alias}
        else:
            model_aliases.pop(full_key, None)

        self.save(config)
        return config

    def add_model(self, provider_key: str, model_entry: dict[str, Any]) -> dict[str, Any]:
        config = self.load()
        providers = self._providers_section(config)
        provider = providers.setdefault(provider_key, {})
        models = provider.setdefault("models", [])
        model_id = str(model_entry.get("id", ""))
        existing = next((item for item in models if str(item.get("id", "")) == model_id), None)
        if existing is None:
            models.append(model_entry)
        else:
            existing.clear()
            existing.update(model_entry)
        self.save(config)
        return config

    def remove_model(self, full_key: str) -> dict[str, Any]:
        provider_key, model_id = self._split_model_key(full_key)
        config = self.load()
        providers = self._providers_section(config)
        provider = providers.get(provider_key)
        if provider:
            models = provider.setdefault("models", [])
            provider["models"] = [item for item in models if str(item.get("id", "")) != model_id]
        self._aliases_section(config).pop(full_key, None)
        self.save(config)
        return config

    def ensure_copy(self, destination: Path) -> None:
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(self.config_path, destination)

    def _providers_section(self, config: dict[str, Any]) -> dict[str, Any]:
        return config.setdefault("models", {}).setdefault("providers", {})

    def _aliases_section(self, config: dict[str, Any]) -> dict[str, Any]:
        return config.setdefault("agents", {}).setdefault("defaults", {}).setdefault("models", {})

    def _split_model_key(self, full_key: str) -> tuple[str, str]:
        if "/" not in full_key:
            raise ValueError("model key must be provider/model")
        provider_key, model_id = full_key.split("/", 1)
        return provider_key, model_id


def iter_configured_model_keys(config: dict[str, Any]) -> Iterable[str]:
    repo = OpenClawConfigRepository(Path("/dev/null"))
    for model in repo.configured_models(config):
        yield model.full_key