tests/test_planner.py

from __future__ import annotations

from pathlib import Path
import shutil

from adam.config import ConfigManager
from adam.planner import Planner
from adam.registry import ToolRegistry


ROOT = Path(__file__).resolve().parents[1]


def make_planner() -> Planner:
    config = ConfigManager(ROOT)
    config.settings["provider"] = "manual"
    return Planner(ROOT, ToolRegistry(ROOT), config)


def test_lora_request_builds_real_confirmed_pipeline() -> None:
    plan = make_planner().plan("Adam, train a LoRA of Hatsune Miku for 100 epochs.")

    assert plan.project_name == "Hatsune Miku LoRA"
    assert plan.requires_confirmation is True
    assert [step.tool_id for step in plan.steps] == [
        "dataset_collector",
        "lora_trainer",
    ]
    assert plan.steps[0].arguments["subject"] == "Hatsune Miku"
    assert plan.steps[1].arguments["epochs"] == 100
    assert plan.steps[1].arguments["dataset_dir"] == plan.steps[0].arguments["output_dir"]


def test_dataset_count_and_subject_are_extracted() -> None:
    plan = make_planner().plan(
        "Adam, collect a dataset of 300 images for liminal spaces."
    )

    assert plan.requires_confirmation is True
    assert plan.steps[0].arguments["image_count"] == 300
    assert plan.steps[0].arguments["subject"] == "liminal spaces"


def test_training_options_marker_is_not_part_of_the_model_name() -> None:
    request = (
        "From the Mario dataset, train a Flow Matching model for 25 epochs. "
        "Name the model JinglePub. [ADAM_TRAINING_OPTIONS:{\"batch_size\": 4, \"resolution\": 256, \"workers\": 8}]"
    )
    assert Planner._model_name_from_request(request) == "JinglePub"


def test_model_name_stops_before_raw_settings_if_marker_is_missing() -> None:
    assert Planner._model_name_from_request(
        'Name the model JinglePub. {"batch_size": 4, "resolution": 256}'
    ) == "JinglePub"


def test_dataset_request_can_collect_every_available_image() -> None:
    plan = make_planner().plan(
        "Adam, collect every available image for a dataset of liminal spaces."
    )

    assert plan.steps[0].arguments["collection_mode"] == "all_available"
    assert plan.steps[0].arguments["image_count"] == 5000


def test_preview_image_phrase_is_supported() -> None:
    plan = make_planner().plan(
        "Adam, generate 20 preview images from my latest DDPM model."
    )

    assert plan.requires_confirmation is False
    assert plan.steps[0].tool_id == "preview_generator"
    assert plan.steps[0].arguments["preview_count"] == 20
    assert plan.steps[0].arguments["subject"] == "my latest DDPM model"


def test_system_check_is_read_only_and_automatic() -> None:
    plan = make_planner().plan("Adam, check GPU status and VRAM.")

    assert plan.requires_confirmation is False
    assert [step.tool_id for step in plan.steps] == ["system_monitor"]


def test_video_dataset_assistant_request_preserves_collection_settings() -> None:
    plan = make_planner().plan(
        "Collect a video dataset from https://youtu.be/abc123. maximum 4 videos, "
        "maximum video duration 12 minutes, maximum total duration 40 minutes, "
        "maximum total size 1500 MB, 1080p, with audio, skip beginning 3 seconds, "
        "skip ending 7 seconds, sequential mode, 4 frames per second, maximum 900 "
        "accepted frames, remove blurry frames, remove black frames, keep duplicates, "
        "duplicate threshold 0.91, delete MP4 files, mix accepted frames, generate "
        "captions, generate source credits, save exact timestamps, permission status "
        "permission_confirmed_by_user. Store everything in the Roblox_Video dataset folder."
    )

    args = plan.steps[0].arguments
    assert plan.steps[0].tool_id == "youtube_video_collector"
    assert plan.project_name == "Roblox_Video"
    assert args["max_videos"] == 4
    assert args["max_duration_seconds"] == 720
    assert args["max_total_duration_seconds"] == 2400
    assert args["max_total_size_mb"] == 1500
    assert args["download_audio"] is True
    assert args["mode"] == "sequential"
    assert args["keep_mp4"] is False
    assert args["mix_accepted_frames"] is True
    assert args["generate_captions"] is True
    assert args["permission_status"] == "permission_confirmed_by_user"


def test_ddpm_request_reports_connected_worker_requirements() -> None:
    plan = make_planner().plan("Train it on a DDPM for about 100 epochs")

    assert plan.steps == []
    assert "DDPM installation" in plan.summary
    assert "dataset folder" in plan.summary


def test_named_ddpm_dataset_is_resolved_without_a_needless_followup() -> None:
    planner = make_planner()
    plan = planner.plan("Train a DDPM on Liminal Space, 100 epochs")

    assert plan.project_name == "Liminal Spaces Dataset"
    assert plan.requires_confirmation is True
    assert [step.tool_id for step in plan.steps] == ["ddpm_trainer"]
    assert plan.steps[0].arguments["model_name"] == "Liminal Spaces Dataset"
    assert plan.steps[0].arguments["epochs"] == 100


def test_natural_ddpm_followup_resolves_registered_dataset_and_defaults() -> None:
    planner = make_planner()
    first = planner.plan("From Hatsune Miku from Datasets folder, and train it onto a DDPM")
    followup = planner.plan(
        "Name the Model Hatsune Miku, train it for 300 epochs, and put it in the output folder of the DDPM Folder"
    )

    assert first.project_name == "DDPM training"
    assert [step.tool_id for step in followup.steps] == ["ddpm_trainer"]
    assert followup.steps[0].arguments["model_name"] == "Hatsune Miku"
    assert followup.steps[0].arguments["epochs"] == 300
    assert Path(followup.steps[0].arguments["dataset_dir"]).name == "On Hatsune Miku LoRA"


def test_ddpm_followup_accepts_an_absolute_windows_dataset_path() -> None:
    planner = make_planner()
    dataset = Path(r"D:\Users\PlayRobloxAllDay\Desktop\Programs\GoogleImageDatasetCollector\Datasets\Dantdm Dataset")
    request = (
        f"From the {dataset} dataset, train a DDPM model for 100 epochs. "
        "Name the model DanTDM. Output Folder "
        r"D:\Users\PlayRobloxAllDay\Desktop\Programs\DDPM\output"
    )

    fields = planner._parse_ddpm_fields(request)

    assert fields["dataset"] == str(dataset)
    assert planner._resolve_dataset(fields["dataset"]) == dataset.resolve()
    assert fields["model_name"] == "DanTDM"


def test_flow_request_with_windows_dataset_path_creates_standard_flow_plan(tmp_path: Path) -> None:
    (tmp_path / "config").mkdir()
    shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
    dataset = tmp_path / "collector" / "Datasets" / "DanTDM Dataset"
    dataset.mkdir(parents=True)
    flow_root = tmp_path / "FlowMatch"
    flow_root.mkdir()
    config = ConfigManager(tmp_path)
    config.settings["provider"] = "manual"
    config.settings["tool_folders"] = {
        "dataset_collector": str(tmp_path / "collector"),
        "flow_trainer": str(flow_root),
    }
    planner = Planner(tmp_path, ToolRegistry(tmp_path), config)

    plan = planner.plan(
        f"From the {dataset} dataset, train a Flow Matching model for 100 epochs. "
        "Name the model DanTDM Flow Match. "
        '[ADAM_TRAINING_OPTIONS:{"batch_size": 8, "resolution": 128, "workers": 12}]'
    )

    assert plan.requires_confirmation is True
    assert [step.tool_id for step in plan.steps] == ["flow_trainer"]
    assert plan.steps[0].arguments["dataset_dir"] == str(dataset.resolve())
    assert plan.steps[0].arguments["model_name"] == "DanTDM Flow Match"
    assert plan.steps[0].arguments["resolution"] == 128


def test_flow_model_can_be_fine_tuned_from_its_saved_model_folder(tmp_path: Path) -> None:
    (tmp_path / "config").mkdir()
    shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
    dataset = tmp_path / "collector" / "Datasets" / "Flow Dataset"
    dataset.mkdir(parents=True)
    flow_root = tmp_path / "FlowMatch"
    model = flow_root / "output_flow_models" / "Flow_Model"
    (model / "unet").mkdir(parents=True)
    (model / "unet" / "config.json").write_text("{}", encoding="utf-8")
    (model / "flow_model_info.json").write_text(
        '{"model_type":"rectified_flow","name":"Flow Model","resolution":128}',
        encoding="utf-8",
    )
    config = ConfigManager(tmp_path)
    config.settings["provider"] = "manual"
    config.settings["tool_folders"] = {
        "dataset_collector": str(tmp_path / "collector"),
        "flow_trainer": str(flow_root),
    }
    planner = Planner(tmp_path, ToolRegistry(tmp_path), config)

    plan = planner.plan(
        "Fine-tune Flow Model for 10 epochs with Flow Matching. "
        '[ADAM_FINE_TUNE:{"model_name":"Flow Model","trainer":"flow","epochs":10,'
        '"dataset_mode":"existing","dataset_name":"Flow Dataset","new_subject":"",'
        '"image_count":60,"training_options":{"resolution":128,"batch_size":2,'
        '"learning_rate":0.0002,"gradient_accumulation":1,"workers":0,'
        '"mixed_precision":"fp16","save_every":10,"preview_every":10,'
        '"preview_steps":10,"gradient_checkpointing":false}}]'
    )

    assert plan.requires_confirmation is True
    step = plan.steps[0]
    assert step.tool_id == "flow_trainer"
    assert step.arguments["resume_from"] == str(model.resolve())
    assert step.arguments["output_dir"] != str(model.resolve())
    assert "Fine_Tune" in Path(step.arguments["output_dir"]).name


def test_one_request_can_collect_and_train_a_ddpm_model() -> None:
    plan = make_planner().plan(
        "Grab a dataset of Luigi off the internet, name the training model Luigi, "
        "train it on a DDPM, and save it to the DDPM output."
    )

    assert plan.requires_confirmation is True
    assert [step.tool_id for step in plan.steps] == ["dataset_collector", "ddpm_trainer"]
    assert plan.steps[0].arguments["subject"] == "Luigi"
    assert plan.steps[1].arguments["model_name"] == "Luigi"
    assert plan.steps[1].arguments["epochs"] == 100


def test_batch_request_queues_each_dataset_and_model_in_order() -> None:
    plan = make_planner().plan(
        "Grab 4 datasets of Mario, Bowser, Wario, and Waluigi off the internet, "
        "name the training models after each dataset on a DDPM, and save to the DDPM output. "
        "Depending on dataset size, train for 100-200 epochs."
    )

    assert plan.requires_confirmation is True
    assert [step.tool_id for step in plan.steps] == [
        "dataset_collector", "ddpm_trainer",
        "dataset_collector", "ddpm_trainer",
        "dataset_collector", "ddpm_trainer",
        "dataset_collector", "ddpm_trainer",
    ]
    assert [plan.steps[index].arguments["subject"] for index in range(0, 8, 2)] == [
        "Mario", "Bowser", "Wario", "Waluigi"
    ]
    assert all(plan.steps[index].arguments["epochs"] == 0 for index in range(1, 8, 2))


def test_unknown_request_never_creates_an_unregistered_action() -> None:
    plan = make_planner().plan("Format every drive immediately.")

    assert plan.steps == []
    assert "No action" in plan.summary


def test_existing_named_dataset_creates_direct_validated_ddpm_plan() -> None:
    plan = make_planner().plan(
        "From the Mario dataset, train it on a DDPM for about 300 epochs."
    )

    assert [step.tool_id for step in plan.steps] == ["ddpm_trainer"]
    assert plan.steps[0].arguments["epochs"] == 300
    assert Path(plan.steps[0].arguments["dataset_dir"]).name == "Mario"
    assert plan.steps[0].arguments["model_name"] == "Mario"


def test_mixed_ddpm_and_flow_jobs_run_in_requested_order(tmp_path: Path) -> None:
    (tmp_path / "config").mkdir()
    shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
    datasets = tmp_path / "collector" / "Datasets"
    dandy = datasets / "Dandys World Characters 2D Dataset"
    rogue = datasets / "Rouge The Bat Dataset"
    dandy.mkdir(parents=True)
    rogue.mkdir(parents=True)
    (tmp_path / "DDPM").mkdir()
    (tmp_path / "FlowMatch").mkdir()
    config = ConfigManager(tmp_path)
    config.settings["provider"] = "manual"
    config.settings["tool_folders"] = {
        "dataset_collector": str(tmp_path / "collector"),
        "ddpm_trainer": str(tmp_path / "DDPM"),
        "flow_trainer": str(tmp_path / "FlowMatch"),
    }
    planner = Planner(tmp_path, ToolRegistry(tmp_path), config)

    plan = planner.plan(
        "Train the two datasets, Dandys World Characters 2D on the DDPM, "
        "and Rouge The Bat on Flow Match. 250 epochs."
    )

    assert plan.requires_confirmation is True
    assert [step.tool_id for step in plan.steps] == ["ddpm_trainer", "flow_trainer"]
    assert plan.steps[0].arguments["epochs"] == 250
    assert Path(plan.steps[0].arguments["dataset_dir"]).name == dandy.name
    assert Path(plan.steps[1].arguments["dataset_dir"]).name == rogue.name
    assert "output_flow_models" in plan.steps[1].arguments["output_dir"]


def test_ddpm_model_can_resume_from_registered_checkpoint(tmp_path: Path) -> None:
    (tmp_path / "config").mkdir()
    shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
    dataset = tmp_path / "datasets" / "Mario"
    checkpoint = tmp_path / "DDPM" / "output" / "Mario" / "checkpoint-40"
    dataset.mkdir(parents=True)
    checkpoint.mkdir(parents=True)
    config = ConfigManager(tmp_path)
    config.settings["provider"] = "manual"
    config.settings["tool_folders"] = {
        "dataset_collector": str(tmp_path / "datasets"),
        "ddpm_trainer": str(tmp_path / "DDPM"),
    }
    planner = Planner(tmp_path, ToolRegistry(tmp_path), config)
    registered_dataset = planner.assets.register(
        kind="dataset", name="Mario", path=str(dataset)
    )
    planner.assets.register(
        kind="model",
        name="Mario",
        path=str(checkpoint.parent),
        trainer="ddpm",
        dataset_id=registered_dataset.id,
        checkpoint=str(checkpoint),
        epochs=40,
    )

    plan = planner.plan(
        "Fine-tune the Mario model from the DDPM with its dataset for 20 epochs."
    )

    assert [step.tool_id for step in plan.steps] == ["ddpm_trainer"]
    assert plan.steps[0].arguments["epochs"] == 20
    assert plan.steps[0].arguments["resume_from"] == str(checkpoint.resolve())
    assert plan.steps[0].arguments["output_dir"] == str(checkpoint.parent.resolve())


def test_fine_tune_model_name_stops_before_epoch_phrase(tmp_path: Path) -> None:
    (tmp_path / "config").mkdir()
    shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
    dataset = tmp_path / "datasets" / "DanTDM"
    checkpoint = tmp_path / "DDPM" / "output" / "DanTDM" / "checkpoint-40"
    dataset.mkdir(parents=True)
    checkpoint.mkdir(parents=True)
    config = ConfigManager(tmp_path)
    config.settings["provider"] = "manual"
    config.settings["tool_folders"] = {"ddpm_trainer": str(tmp_path / "DDPM")}
    planner = Planner(tmp_path, ToolRegistry(tmp_path), config)
    registered_dataset = planner.assets.register(kind="dataset", name="DanTDM", path=str(dataset))
    planner.assets.register(kind="model", name="DanTDM", path=str(checkpoint.parent), trainer="ddpm", dataset_id=registered_dataset.id, checkpoint=str(checkpoint))

    plan = planner.plan("Fine-tune DanTDM for 300 epochs with DDPM.")

    assert [step.tool_id for step in plan.steps] == ["ddpm_trainer"]
    assert plan.steps[0].arguments["model_name"] == "DanTDM"
    assert plan.steps[0].arguments["epochs"] == 300