researchowl/tests/test_shortspec.py

"""Escritura del shot spec: prompt, bucle de reintento y fallback.

El LLM entra como un callable, así que aquí se prueba el bucle, no a Haiku:
respuesta buena, respuesta malformada, typo en una prop, y las tres seguidas.
"""
import json
from pathlib import Path

import pytest

from src.generator.shortspec import (
    MAX_ATTEMPTS, ShortSpecWriter, SpecWriteFailed, extract_json,
)
from tests.test_spec_contract import TEMPLATES

EXAMPLE = Path(__file__).resolve().parents[1] / "src/generator/examples/jal1628.json"

GOOD = {
    "version": 1,
    "meta": {"id": "caso", "title": "Un caso"},
    "shots": [
        {"template": "radar_sweep", "duration": 15.0, "props": {"headline": "3 RADARS"}},
        {"template": "scale_bars", "duration": 15.0, "props": {
            "headline": "REPORTED SCALE",
            "bars": [{"label": "BOEING 747", "value": 232, "unit": "FT"}]}},
    ],
}


class FakeLLM:
    """Devuelve respuestas de una cola y guarda los prompts que recibió."""

    def __init__(self, *responses):
        self.responses = list(responses)
        self.prompts: list[str] = []
        self.systems: list[str] = []

    async def __call__(self, system, prompt):
        self.systems.append(system)
        self.prompts.append(prompt)
        return self.responses.pop(0) if len(self.responses) > 1 else self.responses[0]


def writer(*responses, **kw):
    llm = FakeLLM(*responses)
    return ShortSpecWriter(llm, TEMPLATES, **kw), llm


# --- parseo -----------------------------------------------------------------

def test_extract_json_survives_markdown_fences():
    assert extract_json('```json\n{"a": 1}\n```') == {"a": 1}
    assert extract_json('Here you go:\n{"a": 1}\nHope that helps') == {"a": 1}
    assert extract_json('{"a": 1}') == {"a": 1}


def test_extract_json_complains_when_there_is_no_object():
    with pytest.raises(ValueError):
        extract_json("I'm afraid I can't do that")


def test_straight_quotes_inside_a_string_are_repaired():
    """Cómo falla esto en la vida real (sesión de Bélgica, 2026-08-01): el
    modelo escribe la cita con comillas rectas, que cierran la cadena JSON antes
    de tiempo. Se arregla aquí porque además es lo que se quiere dibujar."""
    broken = '{"quote_a": ""CREDIBLE PEOPLE. THEY TOLD WHAT THEY SAW."", "n": 1}'
    assert extract_json(broken) == {
        "quote_a": "“CREDIBLE PEOPLE. THEY TOLD WHAT THEY SAW.”", "n": 1}


def test_the_repair_leaves_correct_json_alone():
    good = {"a": 'texto con “tipográficas” dentro', "b": [1, 2], "c": {"d": "e"}}
    assert extract_json(json.dumps(good, ensure_ascii=False)) == good


def test_the_repair_does_not_eat_escaped_quotes():
    assert extract_json(r'{"a": "dijo \"hola\" y se fue"}') == {"a": 'dijo "hola" y se fue'}


def test_an_unrepairable_response_reports_the_offending_fragment():
    """El modelo no ve su salida numerada: "line 189 column 22" no le sirve; el
    trozo sí."""
    with pytest.raises(ValueError) as exc:
        extract_json('{"a": 1, "b": [1, 2,,,], "c": 3}')
    assert "aquí:" in str(exc.value)


# --- el prompt --------------------------------------------------------------

def test_prompt_carries_the_fetched_contract_not_a_copy():
    w, _ = writer("{}")
    prompt = w.build_prompt("Caso X", "material", None, "THEEXCLUSIONZONE.COM")
    assert "radar_sweep:" in prompt and "contact_bearing_deg" in prompt
    assert "1-3 elementos" in prompt          # los límites de longitud, del esquema
    assert "ink, amber, amber_dark" in prompt  # la paleta, también del esquema


def test_prompt_states_the_editorial_constraints():
    w, _ = writer("{}")
    prompt = w.build_prompt("Caso X", "material", "https://x.test/post", "X.TEST")
    assert "20-45 seconds" in prompt
    assert "never hex" in prompt
    assert "https://x.test/post" in prompt
    assert "X.TEST" in prompt
    assert "material" in prompt


def test_prompt_includes_the_worked_example_in_full():
    w, _ = writer("{}")
    prompt = w.build_prompt("Caso X", "material", None, "X.TEST")
    example = json.loads(EXAMPLE.read_text())
    assert example["meta"]["title"] in prompt
    assert "counter_close" in prompt


def test_the_worked_example_narrates_most_of_its_shots():
    """El ejemplo es la señal de formato más fuerte del prompt, más que cualquier
    regla en prosa. Sin narración en él, el modelo escribía specs mudos aunque la
    sección 3b le dijera lo contrario: en la primera generación tras añadir la voz
    narró 1 de 8 planos. Si alguien vuelve a dejar el ejemplo mudo, esto salta.
    """
    example = json.loads(EXAMPLE.read_text())
    shots = example["shots"]
    spoken = [s for s in shots if s.get("narration")]

    assert len(spoken) >= len(shots) * 0.6, "el ejemplo enseña a no narrar"

    # Y el silencio del ejemplo es una decisión, no un olvido: calla justo donde
    # la plantilla ya dibuja una cita.
    silent = {s["template"] for s in shots if not s.get("narration")}
    assert silent == {"scale_bars", "document_quote"}


def test_the_worked_example_declares_time_for_its_own_narration():
    """Un plano que se queda corto para su propia voz enseña a infradeclarar: el
    render no corta la voz, alarga el plano, y el total se va del objetivo."""
    from src.generator.spec_contract import NARRATION_WORDS_PER_SECOND, spoken_seconds

    example = json.loads(EXAMPLE.read_text())
    for i, shot in enumerate(example["shots"]):
        narration = shot.get("narration", "")
        if not narration:
            continue
        # La cuenta que el prompt le pide al modelo, aplicada al ejemplo que le
        # pone delante. Si no cuadran, la regla en prosa pierde.
        rule = len(narration.split()) / NARRATION_WORDS_PER_SECOND + 0.5
        assert shot["duration"] >= rule, f"shot {i} declara menos de lo que habla"
        assert shot["duration"] >= spoken_seconds(narration), \
            f"shot {i} se quedaría corto para su propia voz"


def test_the_prompt_gives_a_budget_the_model_can_count():
    """"20-45 segundos" no es accionable: la duración real no está escrita en el
    spec, sale de sumar el mayor entre lo declarado y lo que tarda la voz. El
    modelo sí puede contar sus `duration` y sus palabras, así que el encargo se
    le da en esas dos unidades."""
    from src.generator.shortspec import NARRATION_WORD_BUDGET
    from src.generator.spec_contract import NARRATION_WORDS_PER_SECOND

    w, _ = writer("{}")
    prompt = w.build_prompt("Caso X", "material", None, "X.TEST")

    assert f"{NARRATION_WORDS_PER_SECOND:.1f} words a second" in prompt, \
        "sin el ritmo de la voz no hay cuenta que el modelo pueda hacer"
    assert f"words ÷ {NARRATION_WORDS_PER_SECOND:.1f}" in prompt
    assert f"{NARRATION_WORD_BUDGET} words" in prompt


def test_the_worked_example_obeys_the_budget_it_preaches():
    """El ejemplo es la señal más fuerte del prompt — más que cualquier regla en
    prosa. Uno que hablara de más enseñaría a hablar de más, dijera lo que
    dijera la sección 3b."""
    from src.generator.shortspec import (
        MAX_SHOT_DURATION, NARRATION_WORDS_PER_LINE,
        NARRATION_WORDS_PER_LINE_MAX, NARRATION_WORD_BUDGET,
    )

    example = json.loads(EXAMPLE.read_text())
    lines = [len(s["narration"].split()) for s in example["shots"] if s.get("narration")]

    assert max(s["duration"] for s in example["shots"]) == MAX_SHOT_DURATION

    assert sum(lines) <= NARRATION_WORD_BUDGET
    assert max(lines) <= NARRATION_WORDS_PER_LINE_MAX
    # El tope corto se anuncia como "la media del ejemplo": si deja de serlo, la
    # regla en prosa se convierte en un número inventado y el modelo la nota.
    assert round(sum(lines) / len(lines)) == NARRATION_WORDS_PER_LINE


def test_prompt_says_out_loud_that_there_is_no_article_yet():
    w, _ = writer("{}")
    assert "No article URL yet" in w.build_prompt("X", "m", None, "X.TEST")


def test_prompt_offers_the_live_audio_palette():
    """La mitad de audio del contrato vivo: los presets y sus notas de mood
    vienen de GET /audio, no de este repo."""
    palette = {"sonar": "the case-file mood", "pulse": "tension, built for debunks"}
    w, _ = writer("{}", presets=palette)
    prompt = w.build_prompt("X", "m", None, "X.TEST")
    assert '"pulse" — tension, built for debunks' in prompt
    assert "whose mood fits the shape" in prompt


def test_prompt_without_a_palette_only_offers_the_baseline():
    w, _ = writer("{}")
    prompt = w.build_prompt("X", "m", None, "X.TEST")
    assert '"sonar"' in prompt and '"none"' in prompt
    assert '"pulse"' not in prompt


# --- bucle ------------------------------------------------------------------

@pytest.mark.asyncio
async def test_a_good_response_validates_on_the_first_attempt():
    w, llm = writer(json.dumps(GOOD))
    result = await w.write("Caso X", "material")
    assert result.attempts == 1
    assert result.spec["meta"]["id"] == "caso"
    assert len(llm.prompts) == 1


@pytest.mark.asyncio
async def test_malformed_json_triggers_a_retry():
    w, llm = writer("no soy JSON", json.dumps(GOOD))
    result = await w.write("Caso X", "material")
    assert result.attempts == 2
    assert "no es un objeto JSON válido" in result.history[0][0]
    assert "previous attempt was rejected" in llm.prompts[1]


@pytest.mark.asyncio
async def test_the_exact_error_paths_are_fed_back_verbatim():
    """La ruta que devuelve la validación es lo más útil que se le puede dar al
    modelo: se le pasa tal cual, sin parafrasear."""
    bad = json.loads(json.dumps(GOOD))
    bad["shots"][0]["props"]["sweeeps"] = 2
    w, llm = writer(json.dumps(bad), json.dumps(GOOD))

    result = await w.write("Caso X", "material")

    assert result.attempts == 2
    assert "shots.0.radar_sweep.props.sweeeps" in llm.prompts[1]
    assert "sweeps" in llm.prompts[1]          # y cuáles sí valen


@pytest.mark.asyncio
async def test_three_failures_raise_but_keep_the_last_attempt():
    """La parte cara es la generación, no el render: el último intento viaja en
    la excepción para poder editarlo a mano y reenviarlo."""
    bad = json.loads(json.dumps(GOOD))
    bad["meta"]["id"] = "Caso Con Espacios"
    w, _ = writer(json.dumps(bad))

    with pytest.raises(SpecWriteFailed) as exc:
        await w.write("Caso X", "material")

    assert exc.value.attempts == MAX_ATTEMPTS
    assert exc.value.last_spec["meta"]["id"] == "Caso Con Espacios"
    assert any("meta.id" in e for e in exc.value.errors)


@pytest.mark.asyncio
async def test_an_off_target_duration_is_commented_once_then_accepted():
    """70 s cumple el contrato pero no el encargo: se comenta UNA vez y, si el
    modelo insiste, se renderiza igual antes que tirar la generación.

    Una y no dos. El tercer intento se reserva para el contrato, que sí es
    binario: un spec largo se ve, uno malformado no se puede ni renderizar.
    """
    long_spec = json.loads(json.dumps(GOOD))
    long_spec["shots"][0]["duration"] = 55.0      # 70 s en total
    w, llm = writer(json.dumps(long_spec))

    result = await w.write("Caso X", "material")

    assert result.attempts == 2, "una nota no vale dos reescrituras"
    assert len(llm.prompts) == 2
    assert result.notes and "recorta" in result.notes[0]
    assert "off-brief" in llm.prompts[1]


@pytest.mark.asyncio
async def test_a_valid_attempt_is_not_thrown_away_by_a_worse_one():
    long_spec = json.loads(json.dumps(GOOD))
    long_spec["shots"][0]["duration"] = 55.0
    w, _ = writer(json.dumps(long_spec), "esto ya no es JSON", "tampoco")

    result = await w.write("Caso X", "material")

    assert result.spec["shots"][0]["duration"] == 55.0
    assert result.notes


@pytest.mark.asyncio
async def test_the_rewrite_is_kept_when_it_obeys_the_note_only_halfway():
    """Obedecer a medias es obedecer. Antes se guardaba el PRIMER intento válido
    y se descartaba la reescritura entera, así que un spec que había bajado de
    70 s a 50 s salía a 70."""
    long_spec = json.loads(json.dumps(GOOD))
    long_spec["shots"][0]["duration"] = 55.0      # 70 s
    better = json.loads(json.dumps(GOOD))
    better["shots"][0]["duration"] = 35.0         # 50 s: sigue pasándose, pero menos
    w, _ = writer(json.dumps(long_spec), json.dumps(better))

    result = await w.write("Caso X", "material")

    assert result.spec["shots"][0]["duration"] == 35.0
    assert result.attempts == 2


@pytest.mark.asyncio
async def test_a_rewrite_that_makes_it_worse_is_discarded():
    long_spec = json.loads(json.dumps(GOOD))
    long_spec["shots"][0]["duration"] = 55.0      # 70 s
    worse = json.loads(json.dumps(GOOD))
    worse["shots"][0]["duration"] = 90.0          # 105 s
    w, _ = writer(json.dumps(long_spec), json.dumps(worse))

    result = await w.write("Caso X", "material")

    assert result.spec["shots"][0]["duration"] == 55.0
    assert result.attempts == 2, "se pagaron dos generaciones aunque valga la primera"


@pytest.mark.asyncio
async def test_a_note_does_not_eat_the_attempt_the_contract_needs():
    """Si la reescritura sale malformada, aún queda un intento para arreglarla."""
    long_spec = json.loads(json.dumps(GOOD))
    long_spec["shots"][0]["duration"] = 55.0
    w, llm = writer(json.dumps(long_spec), "esto no es JSON", json.dumps(GOOD))

    result = await w.write("Caso X", "material")

    assert result.attempts == 3 and result.notes == []
    assert "off-brief" in llm.prompts[1]
    assert "not a valid JSON" in llm.prompts[2] or "no es un objeto JSON" in llm.prompts[2]


@pytest.mark.asyncio
async def test_the_contract_is_refetched_after_a_validation_failure():
    """Si el renderizador se actualizó a mitad de la run, la plantilla nueva
    entra en el segundo intento."""
    new_template = {"type": "object", "additionalProperties": False,
                    "required": ["title"],
                    "properties": {"title": {"type": "string", "minLength": 1}}}
    refreshed = {**TEMPLATES, "holo_scan": new_template}

    async def refresh():
        return refreshed

    with_new = json.loads(json.dumps(GOOD))
    with_new["shots"][1] = {"template": "holo_scan", "duration": 15.0,
                            "props": {"title": "X"}}
    w, llm = writer(json.dumps(with_new), json.dumps(with_new),
                    refresh_templates=refresh)

    result = await w.write("Caso X", "material")

    assert result.attempts == 2
    assert "holo_scan" in llm.prompts[1]
    assert result.spec["shots"][1]["template"] == "holo_scan"


@pytest.mark.asyncio
async def test_progress_is_reported_only_when_it_retries():
    seen = []

    async def on_progress(text):
        seen.append(text)

    w, _ = writer("no JSON", json.dumps(GOOD))
    await w.write("Caso X", "material", on_progress=on_progress)
    assert len(seen) == 1 and "attempt 2/3" in seen[0]