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

from __future__ import annotations

import json
from types import SimpleNamespace

from langchain_core.messages import AIMessage, ToolMessage
from obagent_sdk import observe

from production_build_agents.observability import global_data


class _Capture:
    def __init__(self) -> None:
        self.llm: list[dict] = []
        self.tools: list[dict] = []

    def record_llm(self, **payload) -> None:
        self.llm.append(payload)

    def record_tool(self, name, args=None, result=None, **payload) -> None:
        self.tools.append(
            {
                "name": name,
                "args": args,
                "result": result,
                **payload,
            }
        )


def test_agent_trace_reports_metrics_without_raw_messages(
    monkeypatch,
) -> None:
    capture = _Capture()
    owner = SimpleNamespace(
        active=True,
        warn_once=lambda *_args, **_kwargs: None,
    )
    token = global_data._ACTIVE.set(owner)
    monkeypatch.setattr(observe, "current", lambda: capture)
    messages = [
        AIMessage(
            content="客户秘密正文",
            tool_calls=[
                {
                    "id": "call-1",
                    "name": "inspect_tool",
                    "args": {
                        "url": "https://secret.example/?sig=private",
                        "query": "客户秘密查询",
                    },
                }
            ],
            usage_metadata={
                "input_tokens": 10,
                "output_tokens": 5,
                "total_tokens": 15,
            },
        ),
        ToolMessage(
            content=json.dumps(
                {
                    "success": True,
                    "secret_result": "客户秘密工具结果",
                    "_duration_ms": 20,
                    "_operation_id": "Task1-executor-v1:1",
                },
                ensure_ascii=False,
            ),
            tool_call_id="call-1",
            name="inspect_tool",
        ),
    ]
    try:
        global_data.record_sanitized_agent_messages(
            role="executor",
            agent_run_id="safe-agent-run",
            model=SimpleNamespace(model_name="safe-model"),
            messages=messages,
            message_indexes={0, 1},
        )
    finally:
        global_data._ACTIVE.reset(token)

    encoded = json.dumps(
        {"llm": capture.llm, "tools": capture.tools},
        ensure_ascii=False,
    )
    assert capture.llm[0]["input_tokens"] == 10
    assert capture.llm[0]["output_tokens"] == 5
    assert capture.tools[0]["name"] == "inspect_tool"
    assert capture.tools[0]["result"]["operation_id"] == (
        "Task1-executor-v1:1"
    )
    assert "客户秘密正文" not in encoded
    assert "secret.example" not in encoded
    assert "客户秘密查询" not in encoded
    assert "客户秘密工具结果" not in encoded