tests/test_intelligence.py

from datetime import datetime, timezone
from pathlib import Path

from adam.experiment_tracker import ExperimentRun
from adam.generations import GenerationRecord
from adam.intelligence import build_model_intelligence, recommended_training_request


def _run(*, model_name: str = "Atlas", loss: float | None = 0.2, epochs: int = 20) -> ExperimentRun:
    return ExperimentRun(
        id="EXP-1", job_id="1", timestamp=datetime.now(timezone.utc).isoformat(),
        model_architecture="ddpm", model_name=model_name, trigger_word="", base_model="",
        dataset_path="D:/dataset", dataset_name="dataset", dataset_item_count=24,
        epochs=epochs, batch_size=2, learning_rate=0.0001, optimizer="adamw",
        scheduler="", resolution=256, seed=7, status="finished", training_time_seconds=30,
        final_loss=loss, output_folder="D:/models/atlas", checkpoint_paths=[], preview_images=[],
        peak_vram_gb=None, hardware={}, settings={"batch_size": 2}, generation_settings={},
    )


def test_intelligence_requests_visual_review_without_generation_evidence() -> None:
    profile = build_model_intelligence([_run()], [])[0]

    assert profile.state == "Needs visual review"
    assert profile.recommended_epochs == 20


def test_intelligence_detects_plateau_and_creates_follow_up_request(tmp_path: Path) -> None:
    image = tmp_path / "sample.png"
    image.write_bytes(b"not decoded by this test")
    newest = _run(loss=0.200, epochs=20)
    older = _run(loss=0.201, epochs=20)
    older.id = "EXP-0"  # type: ignore[misc]
    record = GenerationRecord(
        metadata_path=tmp_path / "generation.json", folder=tmp_path, images=(image,),
        provider_id="ddpm_generator", provider_name="DDPM", model_name="Atlas",
        model_path="", prompt="test", seed=1, steps=10, sampler="ddim", aspect_ratio="1:1",
        created_at=datetime.now(timezone.utc).isoformat(), smart_generation={}, image_evaluations={
            str(image.resolve()): {"score": 0.8},
        },
    )

    profile = build_model_intelligence([newest, older], [record], root=tmp_path)[0]

    assert profile.state == "Likely plateau"
    assert profile.recommended_epochs == 15
    request = recommended_training_request(profile)
    assert "15 epochs" in request
    assert "Atlas Follow-up" in request