video-image-production-build/tests/test_evidence.py

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace

from langchain_core.tools import tool
from PIL import Image

from production_build_agents.evidence import (
    ValidatorEvidenceError,
    invoke_evidence_tool,
    view_image_evidence,
)
from production_build_agents.agents.validator.stage_context import (
    _stage_media_evidence,
)
from production_build_agents.agents.validator.task_context import (
    _media_evidence,
)
from production_build_agents.contracts.models import Artifact


def _artifact(artifact_type: str, uri: str, *, ordinal: int = 1) -> Artifact:
    return Artifact(
        artifact_id=f"Task1-v1-artifact-{ordinal}",
        artifact_type=artifact_type,
        uri=uri,
        content_sha256="0" * 64,
        size_bytes=1,
        description="媒体测试产物",
    )


class _MediaTools:
    def __init__(self, *, omit_last_view: bool = False) -> None:
        self.calls: list[tuple[str, object]] = []

        @tool("probe_media")
        def probe_media(source: str) -> dict:
            """返回测试媒体类型。"""
            self.calls.append(("probe_media", source))
            media_type = {
                ".mp4": "video",
                ".mp3": "audio",
            }.get(Path(source).suffix, "image")
            return {
                "success": True,
                "source": source,
                "media_type": media_type,
                "_duration_ms": 1,
            }

        @tool("extract_frames")
        def extract_frames(video: str, num_frames: int = 5) -> dict:
            """返回一个测试视频帧。"""
            self.calls.append(("extract_frames", (video, num_frames)))
            return {
                "success": True,
                "frames": [
                    {"timestamp_sec": 0.5, "url": f"{video}#frame"}
                ],
            }

        @tool("view_images")
        def view_images(image_sources: list[str]) -> dict:
            """返回模型可见的测试图片。"""
            self.calls.append(("view_images", list(image_sources)))
            selected = image_sources[:-1] if omit_last_view else image_sources
            return {
                "success": True,
                "images": [
                    {
                        "source": source,
                        "data_url": "data:image/png;base64,aGVsbG8=",
                    }
                    for source in selected
                ],
            }

        self.by_name = {
            item.name: item for item in (probe_media, extract_frames, view_images)
        }


class ValidatorMediaEvidenceTest(unittest.TestCase):
    def test_task_and_stage_keep_distinct_video_payloads(self) -> None:
        video = "https://example.test/final.mp4"
        artifact = _artifact("video", video)
        task_tools = _MediaTools()
        task_evidence, task_blocks = _media_evidence(
            skill=SimpleNamespace(
                skill_id="video-validation",
                allowed_tools=(),
            ),
            delivery=SimpleNamespace(artifacts=[artifact]),
            tools=[
                task_tools.by_name["extract_frames"],
                task_tools.by_name["view_images"],
            ],
        )
        self.assertEqual(task_evidence[0]["video"], video)
        self.assertNotIn("artifact_id", task_evidence[0])
        self.assertEqual(task_tools.calls[0], ("extract_frames", (video, 5)))
        self.assertEqual(len(task_blocks), 1)

        stage_tools = _MediaTools()
        stage_evidence, stage_blocks = _stage_media_evidence(
            SimpleNamespace(artifacts=[artifact]),
            list(stage_tools.by_name.values()),
        )
        frame = next(item for item in stage_evidence if item["kind"] == "video_frame")
        self.assertEqual(frame["artifact_id"], artifact.artifact_id)
        self.assertNotIn("video", frame)
        probe = next(item for item in stage_evidence if item["kind"] == "media_probe")
        self.assertNotIn("_duration_ms", probe)
        self.assertEqual(len(stage_blocks), 1)

    def test_reference_task_probes_each_artifact_before_viewing(self) -> None:
        tools = _MediaTools()
        video = _artifact("video", "https://example.test/reference.mp4")
        audio = _artifact(
            "audio",
            "https://example.test/reference.mp3",
            ordinal=2,
        )
        evidence, blocks = _media_evidence(
            skill=SimpleNamespace(
                skill_id="reference-validation",
                allowed_tools=(),
            ),
            delivery=SimpleNamespace(artifacts=[video, audio]),
            tools=list(tools.by_name.values()),
        )
        probes = [item for item in evidence if item["kind"] == "media_probe"]
        frame = next(item for item in evidence if item["kind"] == "video_frame")
        self.assertEqual(
            [item["artifact_id"] for item in probes],
            [video.artifact_id, audio.artifact_id],
        )
        self.assertEqual(frame["artifact_id"], video.artifact_id)
        self.assertEqual(len(blocks), 1)

    def test_images_are_viewed_in_batches_of_twelve(self) -> None:
        tools = _MediaTools()
        sources = [f"https://example.test/{index}.png" for index in range(13)]
        blocks = view_image_evidence(
            tools.by_name["view_images"],
            sources,
            incomplete_prefix="缺少图片",
        )
        batches = [value for name, value in tools.calls if name == "view_images"]
        self.assertEqual([len(batch) for batch in batches], [12, 1])
        self.assertEqual(len(blocks), 13)

    def test_local_path_is_converted_and_every_source_is_required(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            image_path = Path(temp_dir) / "frame.png"
            Image.new("RGB", (1, 1)).save(image_path)

            @tool("view_images")
            def local_view(image_sources: list[str]) -> dict:
                """只返回本地路径。"""
                return {
                    "success": True,
                    "images": [
                        {
                            "source": image_sources[0],
                            "local_path": str(image_path),
                        }
                    ],
                }

            blocks = view_image_evidence(
                local_view,
                ["frame-source"],
                incomplete_prefix="缺少图片",
            )
        self.assertTrue(
            blocks[0]["image_url"]["url"].startswith("data:image/png;base64,")
        )

        missing = _MediaTools(omit_last_view=True)
        with self.assertRaisesRegex(ValidatorEvidenceError, "second.png"):
            view_image_evidence(
                missing.by_name["view_images"],
                ["first.png", "second.png"],
                incomplete_prefix="缺少图片",
            )

    def test_tool_failures_are_normalized(self) -> None:
        @tool("probe_media")
        def failed_probe(source: str) -> dict:
            """模拟失败。"""
            return {"success": False, "error": f"无法读取 {source}"}

        with self.assertRaisesRegex(
            ValidatorEvidenceError,
            "probe_media 未返回有效证据",
        ):
            invoke_evidence_tool(failed_probe, {"source": "broken.png"})