video-image-production-build/tests/observability/test_context_report.py

from __future__ import annotations

import json
from pathlib import Path

import httpx
from langchain.agents import create_agent
from langchain_core.language_models.fake_chat_models import (
    FakeMessagesListChatModel,
)
from langchain_core.messages import AIMessage

from production_build_agents.observability.context_report import (
    ObAgentContextClient,
    build_context_report,
    local_checkpoint_ai_counts,
    render_text_report,
)
from production_build_agents.run.langgraph_checkpointer import (
    create_sqlite_checkpointer,
)


def _start(call_id: str, logical_key: str, index: int = 1) -> dict:
    return {
        "type": "note",
        "payload": {
            "event_kind": "model_context_start",
            "physical_call_id": call_id,
            "logical_call_index": index,
            "logical_call_key": logical_key,
            "identity": {
                "agent_run_id": "run:Segment1:v1:executor",
                "role": "segment_executor",
                "segment_id": "Segment1",
                "plan_version": 1,
            },
            "input": {
                "message_count": 3,
                "serialized_bytes": 900,
                "tool_messages": {"serialized_bytes": 200},
                "images": {"count": 1, "decoded_bytes": 300},
            },
            "previous": None,
        },
    }


def _terminal(call_id: str, logical_key: str, index: int = 1) -> dict:
    return {
        "type": "llm",
        "payload": {
            "context_metrics": {
                "event_kind": "model_context_terminal",
                "physical_call_id": call_id,
                "logical_call_index": index,
                "logical_call_key": logical_key,
                "identity": {
                    "agent_run_id": "run:Segment1:v1:executor",
                    "role": "segment_executor",
                    "segment_id": "Segment1",
                    "plan_version": 1,
                },
                "status": "success",
                "usage": {
                    "input_tokens": 100,
                    "output_tokens": 10,
                    "input_tokens_delta": None,
                },
                "transport": {
                    "observed": True,
                    "attempt_count": 2,
                    "retry_count": 1,
                    "attempts": [],
                },
            }
        },
    }


def test_report_groups_calls_and_flags_incomplete_and_duplicates() -> None:
    observed = [
        {
            "run": {
                "id": 10,
                "agent": "production",
                "status": "failed",
                "meta": {
                    "thread_id": "round-production",
                    "pipeline_round_id": "round",
                },
                "input_tokens": 200,
                "output_tokens": 20,
            },
            "flat": {
                "instances": [
                    {
                        "flow": [
                            _start("physical-1", "logical-1"),
                            _terminal("physical-1", "logical-1"),
                            _start("physical-2", "logical-1"),
                        ]
                    }
                ]
            },
            "message_call_counts": {"instance-1": 1},
        }
    ]

    report = build_context_report(observed)

    assert report["totals"] == {
        "agent_runs": 1,
        "physical_calls": 2,
        "terminal_calls": 1,
        "incomplete_calls": 1,
        "raw_message_calls": 1,
        "anomalies": 3,
    }
    agent = report["agent_runs"][0]
    assert agent["transport_retry_count"] == 1
    assert agent["max_input_tokens"] == 100
    assert agent["max_context_bytes"] == 900
    kinds = [item["kind"] for item in report["anomalies"]]
    assert kinds.count("duplicate_logical_call") == 2
    assert kinds.count("start_without_terminal") == 1
    assert "Segment1" in render_text_report(report)


def test_report_never_introduces_raw_content() -> None:
    report = build_context_report(
        [
            {
                "run": {"id": 1, "meta": {}},
                "flat": {
                    "instances": [
                        {"flow": [_start("physical-1", "logical-1")]}
                    ]
                },
                "message_call_counts": {},
            }
        ]
    )
    serialized = json.dumps(report, ensure_ascii=False)

    assert "https://" not in serialized
    assert "base64" not in serialized
    assert "prompt" not in serialized.lower()


def test_local_checkpoint_counts_latest_ai_messages(tmp_path: Path) -> None:
    run_dir = tmp_path / "run"
    model = FakeMessagesListChatModel(
        responses=[AIMessage(content="done")]
    )
    with create_sqlite_checkpointer(
        run_dir,
        filename="agent_checkpoints.sqlite",
    ) as checkpointer:
        agent = create_agent(
            model=model,
            tools=[],
            system_prompt="system",
            checkpointer=checkpointer,
        )
        agent.invoke(
            {"messages": [{"role": "user", "content": "hello"}]},
            {"configurable": {"thread_id": "agent-1"}},
        )

    assert local_checkpoint_ai_counts([run_dir]) == {"agent-1": 1}


def test_obagent_client_resolves_pipeline_runs_and_reads_message_counts() -> None:
    requests: list[str] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(str(request.url))
        if request.url.path == "/api/runs":
            page = request.url.params.get("page")
            item = {
                "id": 10 if page == "1" else 11,
                "agent": "global_data" if page == "1" else "production",
                "meta": {
                    "thread_id": (
                        "round-global-data"
                        if page == "1"
                        else "round-production"
                    ),
                    "pipeline_round_id": "round",
                },
            }
            return httpx.Response(
                200,
                json={
                    "items": [item],
                    "page": int(page),
                    "page_size": 1,
                    "total": 2,
                },
            )
        if request.url.path == "/api/runs/10/flat":
            return httpx.Response(
                200,
                json={
                    "instances": [
                        {
                            "id": 100,
                            "flow": [{"type": "llm", "payload": {}}],
                        }
                    ]
                },
            )
        if request.url.path == "/api/instances/100/messages":
            return httpx.Response(
                200,
                json={"inst_id": 100, "calls": [{}, {}]},
            )
        raise AssertionError(request.url)

    client = ObAgentContextClient(
        endpoint="https://obagent.example",
        project="project",
    )
    client.client.close()
    client.client = httpx.Client(
        base_url="https://obagent.example",
        transport=httpx.MockTransport(handler),
    )
    try:
        runs = client.resolve_runs("round")
        observed = client.observed_run(runs[0])
    finally:
        client.close()

    assert [run["id"] for run in runs] == [10, 11]
    assert observed["message_call_counts"] == {"100": 2}
    assert any("page=2" in url for url in requests)