tests/test_llm.py

"""
Test the LLM + parser only. No audio, no Gradio, no Kokoro.

Usage:
    python test_llm.py                          # uses LULLABY_GGUF env var or ./lullaby.gguf
    LULLABY_GGUF=/path/to/x.gguf python test_llm.py
"""

import os
import re
import sys
import time

from llama_cpp import Llama


MODEL_PATH = os.environ.get("LULLABY_GGUF", "lolaby.gguf")

SYSTEM_PROMPT = (
    "You write personalized lullabies for small children, with chord markers "
    "and a tempo/meter header so a guitar accompaniment can be rendered. "
    "Output only the lullaby — no preamble."
)

# Same regexes as app.py — if these match, the audio pipeline will too.
TEMPO_RE = re.compile(r"Tempo:\s*(\d+)\s*bpm,\s*(\d+)/(\d+)", re.IGNORECASE)
PROG_RE = re.compile(r"Progression:\s*([A-G][^\n]+)", re.IGNORECASE)
CHORD_RE = re.compile(r"\[([A-G][^\]]*)\]\s*([^\[\n]*)")


TEST_PROMPTS = [
    (
        "Mia, sleepy and comforted, 6/8",
        "Write a lullaby for: Mia, age 3\n"
        "Loves: her stuffed elephant Pip, the moon\n"
        "Mood: sleepy and comforted\n"
        "Key: C major\n"
        "Meter: 6/8",
    ),
    (
        "Theo with fear, minor key",
        "Write a lullaby for: Theo, age 2\n"
        "Loves: trains and tunnels\n"
        "Fears: the dark\n"
        "Mood: tearful, needs soothing\n"
        "Key: A minor\n"
        "Meter: 6/8",
    ),
    (
        "Beatriz, content, 4/4",
        "Write a lullaby for: Beatriz, age 4\n"
        "Loves: the cat next door, songs about the sea\n"
        "Mood: cosy and content\n"
        "Key: D major\n"
        "Meter: 4/4",
    ),
]


def parse_lullaby(text):
    """Same parser as the app — returns dict or raises ValueError."""
    tempo_match = TEMPO_RE.search(text)
    if not tempo_match:
        raise ValueError("missing 'Tempo:' line")
    prog_match = PROG_RE.search(text)
    if not prog_match:
        raise ValueError("missing 'Progression:' line")

    progression = [c.strip() for c in re.split(r"\s*-\s*", prog_match.group(1)) if c.strip()]

    lines = []
    for raw in text.splitlines():
        raw = raw.strip()
        if not raw.startswith("["):
            continue
        fragments = CHORD_RE.findall(raw)
        if fragments:
            lines.append(fragments)

    if not lines:
        raise ValueError("no chord-marked lyric lines")

    return {
        "bpm": int(tempo_match.group(1)),
        "meter": f"{tempo_match.group(2)}/{tempo_match.group(3)}",
        "progression": progression,
        "lines": lines,
    }


def check(label, cond, detail=""):
    mark = "✓" if cond else "✗"
    print(f"    {mark} {label}" + (f" — {detail}" if detail else ""))
    return cond


def main():
    if not os.path.exists(MODEL_PATH):
        print(f"ERROR: model not found at {MODEL_PATH}")
        print("Set LULLABY_GGUF env var or put lullaby.gguf in cwd.")
        sys.exit(1)

    print(f"Loading {MODEL_PATH}...")
    t0 = time.time()
    llm = Llama(
        model_path=MODEL_PATH,
        n_ctx=1024,
        n_threads=4,
        n_gpu_layers=0,
        chat_format="llama-3",
        verbose=False,
    )
    print(f"Loaded in {time.time() - t0:.1f}s\n")

    total_pass = 0
    total_checks = 0

    for label, prompt in TEST_PROMPTS:
        print(f"=== {label} ===")
        print(f"Prompt:\n{prompt}\n")
        t0 = time.time()
        resp = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            temperature=0.75,
            max_tokens=400,
            top_p=0.9,
        )
        gen_time = time.time() - t0
        output = resp["choices"][0]["message"]["content"].strip()
        print(f"Output ({gen_time:.1f}s):\n{output}\n")

        # Parse it
        try:
            parsed = parse_lullaby(output)
            total_pass += check("parses cleanly", True)
            total_pass += check("tempo found", "bpm" in parsed, f"{parsed['bpm']} bpm")
            total_pass += check("progression has ≥3 chords",
                                len(parsed["progression"]) >= 3,
                                str(parsed["progression"]))
            total_pass += check("≥6 chord-marked lines",
                                len(parsed["lines"]) >= 6,
                                f"{len(parsed['lines'])} lines")

            # Name check
            name = prompt.split("for:")[1].split(",")[0].strip()
            found = name.lower() in output.lower()
            total_pass += check(f"name '{name}' in output", found)
            total_checks += 5
        except ValueError as e:
            check("parses cleanly", False, str(e))
            total_checks += 5

        print()

    print(f"\n=== {total_pass}/{total_checks} checks passed ===")
    sys.exit(0 if total_pass == total_checks else 1)


if __name__ == "__main__":
    main()