tools/test_annotate_dmhy_dag_agent.py

"""Smoke tests for the independent DMHY DAG annotation runner."""

from __future__ import annotations

import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

import tools.annotate_dmhy_dag_agent as agent


def make_args(tmp: Path | None = None, **overrides) -> agent.Args:
    values = {
        "dag": Path("dag.json"),
        "output": Path("out.jsonl"),
        "patch_output": (tmp / "patch.jsonl") if tmp is not None else Path("patch.jsonl"),
        "units_output": Path("units.jsonl"),
        "limit": 10,
        "max_requests": 10,
        "units_per_request": 10,
        "max_context_units": 10,
        "workers": 8,
        "coverage_mode": agent.COVERAGE_SHARED_ONLY,
        "min_incoming_count": 2,
        "min_reachable_terminals": 2,
        "max_reachable_terminals": 500,
        "example_count": 2,
        "max_records": 20,
        "records_per_terminal": 1,
        "terminal_sink_threshold": 1000,
        "terminal_cluster_size": 1,
        "llm": True,
        "base_url": "http://example.test/v1",
        "api_key": "test-key",
        "model": "gpt-5.4-mini",
        "http_timeout": 1,
        "retries": 2,
        "preserve_i_labels": False,
        "resume": False,
    }
    values.update(overrides)
    return agent.Args(**values)


def make_unit(unit_id: str, terminal_ids: list[str]) -> dict:
    terminals = [
        {
            "_terminal_id": terminal_id,
            "_terminal_index": index,
            "value_examples": [f"Show {terminal_id} 01 [1080P].mkv"],
            "suffix_examples": ["01 [1080P]"],
            "weight": 1,
        }
        for index, terminal_id in enumerate(terminal_ids)
    ]
    return {
        "unit_id": unit_id,
        "node_id": 1,
        "kind": "shared_suffix",
        "incoming_count": 2,
        "reachable_terminals": len(terminal_ids),
        "reachable_weight": len(terminal_ids),
        "terminal_ids": terminal_ids,
        "prefix_examples": ["Show "],
        "digit_skeleton_examples": [],
        "value_examples": [terminal["value_examples"][0] for terminal in terminals],
        "suffix_examples": ["01 [1080P]"],
        "common_edge_labels": ["01"],
        "_terminals": terminals,
    }


def response_for(units: list[dict], wrong_first_id: str | None = None) -> dict:
    annotations = []
    for index, unit in enumerate(units):
        annotations.append(
            {
                "unit_id": wrong_first_id if index == 0 and wrong_first_id is not None else unit["unit_id"],
                "terminal_ids": unit["terminal_ids"],
                "episode_title_suffixes": [],
                "media_suffixes": ["[1080P]"],
                "title_candidates": [],
                "llm_label": "media",
                "notes": "ok",
            }
        )
    return {"output_text": json.dumps({"annotations": annotations}, ensure_ascii=False)}


class AnnotateDmhyDagAgentTests(unittest.TestCase):
    def test_batch_request_schema_contains_annotations_array(self) -> None:
        args = make_args()
        units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"])]
        body = agent.request_variants(units, args)[0]
        schema = body["text"]["format"]["schema"]
        self.assertEqual(body["text"]["format"]["name"], "dmhy_dag_batch_annotation")
        self.assertIn("annotations", schema["required"])
        self.assertEqual(schema["properties"]["annotations"]["items"]["required"][0], "unit_id")
        payload = json.loads(body["input"])
        self.assertEqual([unit["unit_id"] for unit in payload["units"]], ["u1", "u2"])

    def test_batch_mismatch_retries_with_validation_error(self) -> None:
        args = make_args(retries=2)
        units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"])]
        calls: list[dict] = []
        original_post = agent.http_post_json
        original_sleep = agent.time.sleep

        def fake_post(body, _args):
            calls.append(body)
            return response_for(units, wrong_first_id="wrong") if len(calls) == 1 else response_for(units)

        try:
            agent.http_post_json = fake_post
            agent.time.sleep = lambda _seconds: None
            patches = agent.annotate_batch(units, args, use_llm=True)
        finally:
            agent.http_post_json = original_post
            agent.time.sleep = original_sleep

        self.assertEqual(len(calls), 2)
        self.assertIn("Previous attempt failed validation", calls[1]["instructions"])
        self.assertEqual([patch["unit_id"] for patch in patches], ["u1", "u2"])
        self.assertTrue(all(patch["source"] == agent.LLM_SOURCE for patch in patches))
        self.assertTrue(all(not patch["fallback"] for patch in patches))

    def test_resume_skips_existing_completed_units_and_keeps_order(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp = Path(tmpdir)
            args = make_args(tmp, resume=True, units_per_request=10, max_context_units=10, workers=1)
            units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"])]
            args.patch_output.write_text(
                json.dumps(
                    {
                        "unit_id": "u1",
                        "terminal_ids": ["t1"],
                        "episode_title_suffixes": ["Finale"],
                        "media_suffixes": [],
                        "title_candidates": [],
                        "llm_label": "title",
                        "notes": "existing",
                        "source": agent.LLM_SOURCE,
                        "fallback": False,
                        "status": "ok",
                        "attempts": 1,
                        "errors": [],
                    },
                    ensure_ascii=False,
                )
                + "\n",
                encoding="utf-8",
            )
            calls: list[list[str]] = []
            original_post = agent.http_post_json

            def fake_post(body, _args):
                payload = json.loads(body["input"])
                calls.append([unit["unit_id"] for unit in payload["units"]])
                return response_for([units[1]])

            try:
                agent.http_post_json = fake_post
                completed, skipped, llm_requested, pending_unprocessed = agent.annotate_units(units, args)
            finally:
                agent.http_post_json = original_post

            self.assertEqual(skipped, 1)
            self.assertEqual(llm_requested, 1)
            self.assertEqual(pending_unprocessed, 0)
            self.assertEqual(calls, [["u2"]])
            patches = [patch for _unit, patch in completed]
            self.assertEqual([patch["unit_id"] for patch in patches], ["u1", "u2"])
            self.assertEqual(patches[0]["episode_title_suffixes"], ["Finale"])

    def test_max_requests_partial_does_not_fallback_unprocessed_units(self) -> None:
        args = make_args(max_requests=1, units_per_request=2, max_context_units=2, workers=1)
        units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"]), make_unit("u3", ["t3"])]
        calls: list[list[str]] = []
        original_post = agent.http_post_json

        def fake_post(body, _args):
            payload = json.loads(body["input"])
            batch_units = [unit for unit in units if unit["unit_id"] in {item["unit_id"] for item in payload["units"]}]
            calls.append([unit["unit_id"] for unit in payload["units"]])
            return response_for(batch_units)

        try:
            agent.http_post_json = fake_post
            completed, skipped, llm_requested, pending_unprocessed = agent.annotate_units(units, args)
        finally:
            agent.http_post_json = original_post

        self.assertEqual(skipped, 0)
        self.assertEqual(llm_requested, 1)
        self.assertEqual(pending_unprocessed, 1)
        self.assertEqual(calls, [["u1", "u2"]])
        patches = [patch for _unit, patch in completed]
        self.assertEqual([patch["unit_id"] for patch in patches], ["u1", "u2"])
        self.assertTrue(all(not patch["fallback"] for patch in patches))

    def test_resume_continues_with_next_unfinished_batch(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp = Path(tmpdir)
            args = make_args(tmp, resume=True, max_requests=1, units_per_request=1, max_context_units=1, workers=1)
            units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"]), make_unit("u3", ["t3"])]
            args.patch_output.write_text(
                json.dumps(
                    {
                        "unit_id": "u1",
                        "terminal_ids": ["t1"],
                        "episode_title_suffixes": [],
                        "media_suffixes": ["[1080P]"],
                        "title_candidates": [],
                        "llm_label": "media",
                        "notes": "existing",
                        "source": agent.LLM_SOURCE,
                        "fallback": False,
                        "status": "ok",
                        "attempts": 1,
                        "errors": [],
                    },
                    ensure_ascii=False,
                )
                + "\n",
                encoding="utf-8",
            )
            calls: list[list[str]] = []
            original_post = agent.http_post_json

            def fake_post(body, _args):
                payload = json.loads(body["input"])
                calls.append([unit["unit_id"] for unit in payload["units"]])
                return response_for([units[1]])

            try:
                agent.http_post_json = fake_post
                completed, skipped, llm_requested, pending_unprocessed = agent.annotate_units(units, args)
            finally:
                agent.http_post_json = original_post

            self.assertEqual(skipped, 1)
            self.assertEqual(llm_requested, 1)
            self.assertEqual(pending_unprocessed, 1)
            self.assertEqual(calls, [["u2"]])
            self.assertEqual([unit["unit_id"] for unit, _patch in completed], ["u1", "u2"])

    def test_cli_heuristic_sample_skips_terminal_sink_and_writes_dataset_rows(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp = Path(tmpdir)
            dag_path = tmp / "dag.json"
            output_path = tmp / "dmhy_weak.dag_sample.jsonl"
            patch_path = tmp / "patch.jsonl"
            units_path = tmp / "units.jsonl"
            dag_path.write_text(
                json.dumps(
                    {
                        "root": 0,
                        "nodes": [
                            {
                                "id": 0,
                                "children": [{"label": "[A] Test Show - 0", "target": 1}],
                            },
                            {
                                "id": 1,
                                "terminal": False,
                                "children": [{"label": "1 [1080P][WEB-DL]", "target": 2}],
                                "incoming_count": 2,
                                "reachable_terminals": 2,
                                "reachable_weight": 2,
                            },
                            {
                                "id": 2,
                                "terminal": True,
                                "children": [],
                                "incoming_count": 10,
                                "reachable_terminals": 100,
                            },
                        ],
                        "terminals": [
                            {
                                "terminal_id": "t0",
                                "node_id": 1,
                                "prefix": "[A] Test Show - 0",
                                "weight": 1,
                                "suffix_examples": ["1 [1080P][WEB-DL]"],
                                "value_examples": ["[A] Test Show - 01 [1080P][WEB-DL].mkv"],
                            },
                            {
                                "terminal_id": "t1",
                                "node_id": 1,
                                "prefix": "[A] Test Show - 0",
                                "weight": 1,
                                "suffix_examples": ["2 [1080P][WEB-DL]"],
                                "value_examples": ["[A] Test Show - 02 [1080P][WEB-DL].mkv"],
                            },
                        ],
                    },
                    ensure_ascii=False,
                ),
                encoding="utf-8",
            )

            run = subprocess.run(
                [
                    sys.executable,
                    "-m",
                    "tools.annotate_dmhy_dag_agent",
                    "--dag",
                    str(dag_path),
                    "--output",
                    str(output_path),
                    "--patch-output",
                    str(patch_path),
                    "--units-output",
                    str(units_path),
                    "--limit",
                    "1",
                    "--max-requests",
                    "0",
                    "--workers",
                    "2",
                ],
                check=False,
                capture_output=True,
                text=True,
            )
            self.assertEqual(run.returncode, 0, run.stderr)

            patches = [json.loads(line) for line in patch_path.read_text(encoding="utf-8").splitlines()]
            self.assertEqual(len(patches), 1)
            self.assertEqual(patches[0]["unit_id"], "dag-node-1")
            self.assertEqual(patches[0]["terminal_ids"], ["t0", "t1"])

            units = [json.loads(line) for line in units_path.read_text(encoding="utf-8").splitlines()]
            self.assertEqual(len(units), 1)
            self.assertEqual(
                set(units[0]),
                {
                    "unit_id",
                    "source_kind",
                    "source_id",
                    "terminal_ids",
                    "weight",
                    "context",
                    "examples",
                    "expected_output",
                },
            )
            self.assertEqual(units[0]["source_kind"], "prefix_dag")
            self.assertEqual(units[0]["source_id"], "1")
            self.assertEqual(units[0]["terminal_ids"], ["t0", "t1"])
            self.assertEqual(units[0]["weight"], 2)
            self.assertEqual(units[0]["context"]["prefixes"], ["[A] Test Show - 0"])
            self.assertEqual(units[0]["context"]["digit_skeletons"], [])
            self.assertEqual(units[0]["context"]["edge_labels"], ["1 [1080P][WEB-DL]"])
            self.assertIn("incoming_count=2", units[0]["context"]["notes"])
            self.assertEqual(
                units[0]["examples"]["values"],
                [
                    "[A] Test Show - 01 [1080P][WEB-DL].mkv",
                    "[A] Test Show - 02 [1080P][WEB-DL].mkv",
                ],
            )
            self.assertEqual(units[0]["examples"]["suffixes"], ["1 [1080P][WEB-DL]", "2 [1080P][WEB-DL]"])
            self.assertEqual(units[0]["expected_output"], {"schema_version": "dmhy-annotation-v1"})

            rows = [json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines()]
            self.assertEqual(len(rows), 2)
            self.assertEqual(rows[0]["tokens"][0], "[")
            self.assertTrue(all(label == "O" or label.startswith("B-") for row in rows for label in row["labels"]))

    def test_cli_all_terminals_adds_terminal_cluster_for_uncovered_terminal(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp = Path(tmpdir)
            dag_path = tmp / "dag.json"
            output_path = tmp / "dmhy_weak.dag_all.jsonl"
            patch_path = tmp / "patch.jsonl"
            units_path = tmp / "units.jsonl"
            dag_path.write_text(
                json.dumps(
                    {
                        "root": 0,
                        "nodes": [
                            {
                                "id": 0,
                                "children": [
                                    {"label": "A shared", "target": 1},
                                    {"label": "B solo", "target": 4},
                                ],
                            },
                            {
                                "id": 1,
                                "children": [
                                    {"label": " 01", "target": 2},
                                    {"label": " 02", "target": 3},
                                ],
                                "incoming_count": 2,
                                "reachable_terminals": 2,
                                "reachable_weight": 20,
                            },
                            {
                                "id": 2,
                                "children": [],
                                "incoming_count": 1,
                                "reachable_terminals": 1,
                                "reachable_weight": 10,
                            },
                            {
                                "id": 3,
                                "children": [],
                                "incoming_count": 1,
                                "reachable_terminals": 1,
                                "reachable_weight": 10,
                            },
                            {
                                "id": 4,
                                "children": [],
                                "incoming_count": 1,
                                "reachable_terminals": 1,
                                "reachable_weight": 5,
                            },
                        ],
                        "terminals": [
                            {
                                "terminal_id": "t0",
                                "node_id": 2,
                                "prefix": "Shared Show 01",
                                "weight": 10,
                                "suffix_examples": [" [1080P]"],
                                "value_examples": ["Shared Show 01 [1080P].mkv"],
                            },
                            {
                                "terminal_id": "t1",
                                "node_id": 3,
                                "prefix": "Shared Show 02",
                                "weight": 10,
                                "suffix_examples": [" [1080P]"],
                                "value_examples": ["Shared Show 02 [1080P].mkv"],
                            },
                            {
                                "terminal_id": "t2",
                                "node_id": 4,
                                "prefix": "Solo Show 03",
                                "weight": 5,
                                "suffix_examples": [" [WEB-DL]"],
                                "value_examples": ["Solo Show 03 [WEB-DL].mkv"],
                            },
                        ],
                    },
                    ensure_ascii=False,
                ),
                encoding="utf-8",
            )

            run = subprocess.run(
                [
                    sys.executable,
                    "-m",
                    "tools.annotate_dmhy_dag_agent",
                    "--dag",
                    str(dag_path),
                    "--output",
                    str(output_path),
                    "--patch-output",
                    str(patch_path),
                    "--units-output",
                    str(units_path),
                    "--coverage-mode",
                    "all-terminals",
                    "--min-incoming-count",
                    "2",
                    "--min-reachable-terminals",
                    "2",
                    "--limit",
                    "100",
                    "--max-requests",
                    "0",
                ],
                check=False,
                capture_output=True,
                text=True,
            )
            self.assertEqual(run.returncode, 0, run.stderr)
            summary = json.loads(run.stdout)
            self.assertEqual(summary["coverage_mode"], "all-terminals")
            self.assertEqual(summary["dag_terminals"], 3)
            self.assertEqual(summary["unique_terminal_coverage"], 3)

            patches = [json.loads(line) for line in patch_path.read_text(encoding="utf-8").splitlines()]
            self.assertEqual([patch["kind"] for patch in patches], ["shared_suffix", "terminal_cluster"])
            self.assertEqual(patches[0]["terminal_ids"], ["t0", "t1"])
            self.assertEqual(patches[1]["terminal_ids"], ["t2"])

            units = [json.loads(line) for line in units_path.read_text(encoding="utf-8").splitlines()]
            self.assertEqual(len(units), 2)
            self.assertEqual(units[0]["unit_id"], "dag-node-1")
            self.assertEqual(units[0]["terminal_ids"], ["t0", "t1"])
            self.assertEqual(units[1]["unit_id"], "terminal-cluster-1")
            self.assertEqual(units[1]["source_kind"], "prefix_dag")
            self.assertEqual(units[1]["source_id"], "terminal-cluster-1")
            self.assertEqual(units[1]["terminal_ids"], ["t2"])
            self.assertIn("kind=terminal_cluster", units[1]["context"]["notes"])
            self.assertEqual(
                {terminal_id for unit in units for terminal_id in unit["terminal_ids"]},
                {"t0", "t1", "t2"},
            )

            rows = [json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines()]
            self.assertEqual({row["terminal_id"] for row in rows}, {"t0", "t1", "t2"})
            self.assertIn("terminal-cluster-1", {row["unit_id"] for row in rows})


if __name__ == "__main__":
    unittest.main()