video-image-production-build/tests/run/test_metrics.py

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from production_build_agents.run.metrics import (
    record_agent_messages,
    record_phase_duration,
    record_replan_reason,
    record_run_invocation,
)


class _NamedModel:
    model_name = "test-model"


class RunMetricsTest(unittest.TestCase):
    def test_contract_correction_diffs_are_idempotent(self) -> None:
        before = {
            "payload": {"rule": "保持暖咖色"},
            "summary": "保持业务摘要",
            "artifact_binding_claims": [],
            "removed_field": "旧字段",
        }
        after = {
            "payload": {"rule": "保持暖咖色"},
            "summary": "保持业务摘要",
            "artifact_binding_claims": [{"expectation_id": "Expectation1"}],
            "added_field": "新字段",
        }
        media_review = {
            **after,
            "summary": "媒体自检允许修改的摘要",
        }
        messages = [
            AIMessage(id="candidate-1", content=json.dumps(before)),
            HumanMessage(content="上一版未通过运行时校验。错误：Binding 无效。"),
            AIMessage(
                id="candidate-2",
                content="```json\n" + json.dumps(after) + "\n```",
            ),
            HumanMessage(content="最终媒体证据自检：请查看图片。"),
            AIMessage(id="candidate-3", content=json.dumps(media_review)),
        ]

        with tempfile.TemporaryDirectory() as temp_dir:
            run_dir = Path(temp_dir)
            for _ in range(2):
                record_agent_messages(
                    run_dir,
                    run_id="Run-correction-metrics",
                    role="executor",
                    agent_run_id="Run-correction-executor-Task1-v1",
                    model=_NamedModel(),
                    messages=messages,
                )
            metrics = json.loads(
                (run_dir / "run_metrics.json").read_text(encoding="utf-8")
            )

        self.assertEqual(
            metrics["corrections"],
            [
                {
                    "role": "executor",
                    "agent_run_id": "Run-correction-executor-Task1-v1",
                    "attempt": 2,
                    "changed_fields": ["artifact_binding_claims"],
                    "added_fields": ["added_field"],
                    "removed_fields": ["removed_field"],
                }
            ],
        )

    def test_unparseable_correction_does_not_break_metrics(self) -> None:
        messages = [
            AIMessage(id="invalid-1", content="不是 JSON"),
            HumanMessage(content="上一版未通过运行时格式校验。"),
            AIMessage(id="valid-2", content='{"summary": "已修正"}'),
        ]
        with tempfile.TemporaryDirectory() as temp_dir:
            run_dir = Path(temp_dir)
            record_agent_messages(
                run_dir,
                run_id="Run-invalid-correction",
                role="validator",
                agent_run_id="validator-Task1-v1",
                model=_NamedModel(),
                messages=messages,
            )
            metrics = json.loads(
                (run_dir / "run_metrics.json").read_text(encoding="utf-8")
            )

        self.assertEqual(metrics["corrections"], [])

    def test_agent_messages_are_aggregated_idempotently(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            run_dir = Path(temp_dir)
            messages = [
                AIMessage(
                    id="ai-1",
                    content="done",
                    usage_metadata={
                        "input_tokens": 11,
                        "output_tokens": 7,
                        "total_tokens": 18,
                    },
                    response_metadata={
                        "model_name": "reported-model",
                        "cost_usd": 0.0125,
                    },
                ),
                ToolMessage(
                    name="probe_media",
                    tool_call_id="tool-1",
                    content=json.dumps(
                        {
                            "success": True,
                            "_duration_ms": 42,
                            "_operation_replayed": True,
                        }
                    ),
                ),
            ]

            for _ in range(2):
                record_agent_messages(
                    run_dir,
                    run_id="Run-metrics",
                    role="executor",
                    agent_run_id="Run-metrics-executor-Task1-v1",
                    model=_NamedModel(),
                    messages=messages,
                )

            metrics = json.loads(
                (run_dir / "run_metrics.json").read_text(encoding="utf-8")
            )
            self.assertEqual(metrics["model_totals"]["calls"], 1)
            self.assertEqual(metrics["model_totals"]["total_tokens"], 18)
            self.assertEqual(
                metrics["model_totals"]["reported_cost_usd"],
                0.0125,
            )
            self.assertEqual(
                metrics["model_totals"]["cost_status"],
                "reported",
            )
            self.assertEqual(
                metrics["models"]["executor"]["models"],
                ["reported-model"],
            )
            self.assertEqual(metrics["tools"]["calls"], 1)
            self.assertEqual(metrics["tools"]["successful_calls"], 1)
            self.assertEqual(metrics["tools"]["replayed_calls"], 1)
            self.assertEqual(metrics["tools"]["total_duration_ms"], 42)

    def test_phase_invocation_and_replan_reason_are_recorded(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            run_dir = Path(temp_dir)
            record_phase_duration(
                run_dir,
                run_id="Run-metrics",
                phase="execute_task",
                duration_ms=51,
                outcome="RUNNING",
            )
            record_run_invocation(
                run_dir,
                run_id="Run-metrics",
                duration_ms=90,
                status="COMPLETED",
            )
            for _ in range(2):
                record_replan_reason(
                    run_dir,
                    run_id="Run-metrics",
                    plan_version=1,
                    failure_scope="task",
                    details={
                        "task_id": "Task1",
                        "verdict": "FAIL",
                    },
                )

            metrics = json.loads(
                (run_dir / "run_metrics.json").read_text(encoding="utf-8")
            )
            self.assertEqual(
                metrics["phases"]["execute_task"]["total_duration_ms"],
                51,
            )
            self.assertEqual(metrics["run_invocations"]["count"], 1)
            self.assertEqual(
                metrics["run_invocations"]["last_status"],
                "COMPLETED",
            )
            self.assertEqual(len(metrics["replans"]), 1)
            self.assertEqual(
                metrics["replans"][0]["failure_scope"],
                "task",
            )

    def test_stable_phase_and_invocation_events_do_not_double_count(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            run_dir = Path(temp_dir)
            for _ in range(2):
                record_phase_duration(
                    run_dir,
                    run_id="Run-resume-metrics",
                    phase="SEGMENT_VALIDATOR",
                    duration_ms=20,
                    outcome="PASS",
                    event_key="Segment2:v1:validator",
                )
                record_run_invocation(
                    run_dir,
                    run_id="Run-resume-metrics",
                    duration_ms=30,
                    status="COMPLETED",
                    event_key="Segment2:v1:terminal",
                )
            metrics = json.loads(
                (run_dir / "run_metrics.json").read_text(encoding="utf-8")
            )

        self.assertEqual(
            metrics["phases"]["SEGMENT_VALIDATOR"]["calls"],
            1,
        )
        self.assertEqual(metrics["run_invocations"]["count"], 1)


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