adam/config.py

from __future__ import annotations

import json
from copy import deepcopy
from pathlib import Path
from typing import Any


DEFAULT_SETTINGS: dict[str, Any] = {
    "provider": "ollama",
    "ollama_url": "http://localhost:11434",
    "ollama_model": "qwen2.5:1.5b",
    "ollama_chat_max_tokens": 1024,
    "ollama_chat_response_length": "automatic",
    "ollama_proposed_actions": True,
    "web_search_enabled": True,
    "web_link_reading_enabled": True,
    "command_center_mode": "trainer",
    "desktop_notifications": True,
    "sound_notifications": False,
    "ask_before_long_tasks": True,
    "atlas_warning_temperature": 82,
    "atlas_critical_temperature": 90,
    "atlas_warning_disk_gb": 10,
    "atlas_critical_disk_gb": 1,
    "atlas_stall_minutes": 30,
    "trusted_dataset_ddpm_automation": False,
    "demo_step_delay": 0.24,
    "max_dataset_images_without_confirmation": 100,
    "training_presets": {},
    "remote_access": {
        "enabled": False,
        "bind_address": "127.0.0.1",
        "port": 8765,
        "token": "",
        "allow_job_control": False,
        "auto_approve_training": False,
    },
    "tool_folders": {
        "dataset_collector": "",
        "caption_generator": "",
        "lora_trainer": "",
        "ddpm_trainer": "",
        "flow_trainer": "",
        "oasis_trainer": "",
        "wan_video_trainer": "",
        "preview_generator": "",
    },
}


class ConfigManager:
    def __init__(self, root: Path) -> None:
        self.root = root
        self.config_dir = root / "config"
        self.config_dir.mkdir(parents=True, exist_ok=True)
        self.settings_path = self.config_dir / "settings.json"
        self.settings: dict[str, Any] = {}
        self.load()

    def load(self) -> dict[str, Any]:
        self.settings = deepcopy(DEFAULT_SETTINGS)
        if self.settings_path.exists():
            try:
                stored = json.loads(self.settings_path.read_text(encoding="utf-8"))
                if isinstance(stored, dict):
                    self.settings.update(stored)
            except (OSError, json.JSONDecodeError):
                pass
        else:
            self.save()
        return self.settings

    def save(self) -> None:
        temporary = self.settings_path.with_suffix(".tmp")
        temporary.write_text(
            json.dumps(self.settings, indent=2, sort_keys=True),
            encoding="utf-8",
        )
        temporary.replace(self.settings_path)

    def get(self, key: str, default: Any = None) -> Any:
        return self.settings.get(key, default)

    def update(self, values: dict[str, Any]) -> None:
        self.settings.update(values)
        self.save()