code/self_healing_finetuning/model_policy/profiles.py

from __future__ import annotations

import copy
import os
from typing import Any

from n21.config import load_structured
from n21.settings import CONFIG_ROOT


PROFILE_CONFIG = CONFIG_ROOT / "model_profiles.json"


def _profiles_doc() -> dict[str, Any]:
    return load_structured(PROFILE_CONFIG)


def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
    result = copy.deepcopy(base)
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(result.get(key), dict):
            result[key] = _deep_merge(result[key], value)
        else:
            result[key] = copy.deepcopy(value)
    return result


def list_model_profiles() -> dict[str, dict[str, Any]]:
    profiles = _profiles_doc().get("profiles", {})
    if not isinstance(profiles, dict):
        raise ValueError(f"invalid model profiles config: {PROFILE_CONFIG}")
    return profiles


def resolve_model_profile(
    profile_name: str | None = None,
    *,
    model_candidate: str | None = None,
    base_model_id: str | None = None,
) -> dict[str, Any]:
    doc = _profiles_doc()
    profiles = doc.get("profiles", {})
    default_name = str(doc.get("default_profile") or "fingpt")
    name = (profile_name or os.environ.get("SHFT_MODEL_PROFILE") or default_name).strip()
    profile = copy.deepcopy(profiles.get(name))
    if not isinstance(profile, dict):
        raise ValueError(f"unknown SHFT model profile `{name}`; available profiles: {', '.join(sorted(profiles))}")

    env_model = os.environ.get("SHFT_MODEL_CANDIDATE")
    env_base = os.environ.get("SHFT_BASE_MODEL_ID")
    if model_candidate or env_model:
        profile["model_candidate"] = model_candidate or env_model
    if base_model_id or env_base:
        profile["base_model_id"] = base_model_id or env_base
    profile["name"] = name
    profile.setdefault("adapter_bootstrap", False)
    profile.setdefault("baseline_adapter_repo", profile["model_candidate"] if profile["adapter_bootstrap"] else None)
    return profile


def apply_provider_profile(config: dict[str, Any], profile: dict[str, Any]) -> dict[str, Any]:
    merged = _deep_merge(config, profile.get("provider_overrides") or {})
    storage = merged.setdefault("storage", {})
    jobs = merged.setdefault("jobs", {})
    storage["model_repo"] = profile["model_candidate"]
    jobs["base_model_id"] = profile["base_model_id"]
    return merged