85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
"""Model profile loading and resolution utilities."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import yaml
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
DEFAULT_MODEL_CONFIG = REPO_ROOT / "configs" / "models.yaml"
|
|
LEGACY_MODEL_CONFIG = REPO_ROOT / "configs" / "model_profiles.yaml"
|
|
|
|
|
|
class ModelConfigError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def load_model_config(path: Path | None = None) -> dict[str, Any]:
|
|
cfg_path = path or DEFAULT_MODEL_CONFIG
|
|
if not cfg_path.exists() and LEGACY_MODEL_CONFIG.exists():
|
|
cfg_path = LEGACY_MODEL_CONFIG
|
|
if not cfg_path.exists():
|
|
raise ModelConfigError(f"model config not found: {cfg_path}")
|
|
try:
|
|
data = yaml.safe_load(cfg_path.read_text(encoding="utf-8")) or {}
|
|
except Exception as exc:
|
|
raise ModelConfigError(f"invalid YAML in {cfg_path}: {exc}") from exc
|
|
if not isinstance(data, dict):
|
|
raise ModelConfigError(f"invalid model config shape in {cfg_path}")
|
|
return data
|
|
|
|
|
|
def resolve_model_profile(
|
|
*,
|
|
profile: str | None = None,
|
|
overrides: dict[str, str] | None = None,
|
|
path: Path | None = None,
|
|
) -> dict[str, Any]:
|
|
cfg = load_model_config(path)
|
|
profiles = cfg.get("profiles") or {}
|
|
defaults = cfg.get("defaults") or {}
|
|
selected = profile or defaults.get("profile")
|
|
if not selected:
|
|
raise ModelConfigError("no model profile provided and no defaults.profile set")
|
|
if selected not in profiles:
|
|
raise ModelConfigError(f"unknown model profile: {selected}")
|
|
|
|
selected_profile = profiles[selected] or {}
|
|
roles = dict(selected_profile.get("roles") or {})
|
|
task_types = dict(selected_profile.get("task_types") or defaults.get("task_types") or {})
|
|
if defaults.get("script_models"):
|
|
for role, model in (defaults.get("script_models") or {}).items():
|
|
roles.setdefault(role, model)
|
|
for role, model in (overrides or {}).items():
|
|
roles[role] = model
|
|
|
|
return {
|
|
"profile": selected,
|
|
"description": selected_profile.get("description", ""),
|
|
"roles": roles,
|
|
"task_types": task_types,
|
|
}
|
|
|
|
|
|
def list_model_profiles(path: Path | None = None) -> list[str]:
|
|
cfg = load_model_config(path)
|
|
profiles = cfg.get("profiles") or {}
|
|
return sorted(profiles.keys())
|
|
|
|
|
|
def parse_model_overrides(items: list[str] | None) -> dict[str, str]:
|
|
out: dict[str, str] = {}
|
|
for item in items or []:
|
|
if "=" not in item:
|
|
raise ModelConfigError(f"invalid override '{item}', expected role=model")
|
|
role, model = item.split("=", 1)
|
|
role = role.strip()
|
|
model = model.strip()
|
|
if not role or not model:
|
|
raise ModelConfigError(f"invalid override '{item}', expected role=model")
|
|
out[role] = model
|
|
return out
|