test_evaluation.py

"""Tests for check_evaluation behavior."""

from unittest.mock import AsyncMock, MagicMock, patch
from uuid import UUID

import pytest
from agent_control import evaluation
from agent_control.evaluation import EvaluationResult
from agent_control_models import Step
from pydantic import ValidationError


@pytest.mark.asyncio
async def test_check_evaluation_requires_step_name_before_server_call():
    """Typed request validation should reject steps without a name before server call."""

    client = MagicMock()
    client.http_client = AsyncMock()
    client.http_client.post = AsyncMock()

    with pytest.raises(ValidationError):
        await evaluation.check_evaluation(
            client=client,
            agent_name=UUID("00000000-0000-0000-0000-000000000001"),
            step={"type": "llm", "input": "hello"},
            stage="pre",
        )

    client.http_client.post.assert_not_called()


@pytest.mark.asyncio
async def test_check_evaluation_returns_result_model():
    """check_evaluation returns a parsed EvaluationResult."""
    class DummyResponse:
        def raise_for_status(self) -> None:
            return None

        def json(self) -> dict[str, object]:
            return {"is_safe": True, "confidence": 0.75, "reason": "ok"}

    client = MagicMock()
    client.http_client = MagicMock()
    client.http_client.post = AsyncMock(return_value=DummyResponse())

    result = await evaluation.check_evaluation(
        client=client,
        agent_name="Agent-Example_01",
        step={"type": "llm", "name": "chat", "input": "hello"},
        stage="pre",
    )

    assert result.is_safe is True
    assert result.confidence == 0.75
    assert result.reason == "ok"
    client.http_client.post.assert_awaited_once_with(
        "/api/v1/evaluation",
        json={
            "agent_name": "agent-example_01",
            "step": {
                "type": "llm",
                "name": "chat",
                "input": "hello",
                    "output": None,
                    "context": None,
                    "tools": None,
                    "ground_truth": None,
            },
            "stage": "pre",
            "target_type": None,
            "target_id": None,
        },
        headers=None,
    )


@pytest.mark.asyncio
async def test_check_evaluation_sends_custom_step_type_to_server():
    """The server path sends one generic Step inside EvaluationRequest."""

    class DummyResponse:
        def raise_for_status(self) -> None:
            return None

        def json(self) -> dict[str, object]:
            return {"is_safe": True, "confidence": 1.0}

    client = MagicMock()
    client.http_client = MagicMock()
    client.http_client.post = AsyncMock(return_value=DummyResponse())

    await evaluation.check_evaluation(
        client=client,
        agent_name="Agent-Example_01",
        step={"type": "trace", "name": "trace-check", "input": "hello"},
        stage="post",
    )

    request = client.http_client.post.await_args.kwargs["json"]
    assert request["step"] == {
        "type": "trace",
        "name": "trace-check",
        "input": "hello",
        "output": None,
        "context": None,
        "tools": None,
        "ground_truth": None,
    }


@pytest.mark.asyncio
async def test_evaluate_controls_requires_server_url():
    """evaluate_controls should require server_url to be configured."""
    with patch("agent_control.state.server_url", None):
        with pytest.raises(RuntimeError, match="Server URL not configured"):
            await evaluation.evaluate_controls(
                step_name="chat",
                input="hello",
                stage="pre",
                agent_name="test-bot",
            )


@pytest.mark.asyncio
async def test_evaluate_controls_with_explicit_agent_name(monkeypatch):
    """evaluate_controls should call check_evaluation_with_local."""
    mock_result = EvaluationResult(is_safe=True, confidence=1.0)
    mock_check = AsyncMock(return_value=mock_result)
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)

    with patch("agent_control.state.server_url", "http://localhost:8000"):
        with patch("agent_control.state.api_key", None):
            result = await evaluation.evaluate_controls(
                step_name="chat",
                input="hello",
                stage="pre",
                agent_name="test-bot",
            )

    assert result.is_safe is True
    assert result.confidence == 1.0
    mock_check.assert_called_once()


@pytest.mark.asyncio
async def test_evaluate_controls_with_context(monkeypatch):
    """evaluate_controls should pass context through to evaluation."""
    mock_result = EvaluationResult(is_safe=True, confidence=1.0)
    mock_check = AsyncMock(return_value=mock_result)
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)

    with patch("agent_control.state.server_url", "http://localhost:8000"):
        with patch("agent_control.state.api_key", None):
            await evaluation.evaluate_controls(
                step_name="chat",
                input="hello",
                context={"user_id": "123"},
                stage="pre",
                agent_name="test-bot",
            )

    assert mock_check.call_args is not None


@pytest.mark.parametrize("step_type", ["trace", "session"])
@pytest.mark.asyncio
async def test_evaluate_controls_accepts_custom_step_type(monkeypatch, step_type):
    """Direct evaluation preserves custom types in the runtime Step."""
    mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0))
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)

    with patch("agent_control.state.server_url", "http://localhost:8000"):
        await evaluation.evaluate_controls(
            step_name="trace_check",
            step_type=step_type,
            input="trace input",
            output="trace output",
            stage="post",
            agent_name="test-bot",
        )

    step = mock_check.call_args.kwargs["step"]
    assert step.type == step_type
    assert step.name == "trace_check"


@pytest.mark.asyncio
async def test_evaluate_controls_rejects_empty_step_type(monkeypatch):
    """Direct evaluation rejects empty custom types before evaluation."""
    mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0))
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)

    with patch("agent_control.state.server_url", "http://localhost:8000"):
        with pytest.raises(ValueError, match="non-empty string"):
            await evaluation.evaluate_controls(
                step_name="trace_check",
                step_type="",
                agent_name="test-bot",
            )

    mock_check.assert_not_called()


@pytest.mark.asyncio
async def test_evaluate_controls_preserves_explicit_tools_and_ground_truth(monkeypatch):
    """Explicit structured scorer context is preserved on the SDK Step."""
    # Given: a configured SDK and local evaluation boundary
    mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0))
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)
    tools = [
        {"name": "search", "description": "Search", "input_schema": {"type": "object"}}
    ]

    # When: a direct SDK caller supplies tools and ground truth
    with patch("agent_control.state.server_url", "http://localhost:8000"):
        await evaluation.evaluate_controls(
            step_name="chat",
            input="question",
            output="answer",
            tools=tools,
            ground_truth={"answer": "expected"},
            stage="post",
            agent_name="test-bot",
        )

    # Then: both fields are retained as structured JSON
    step = mock_check.call_args.kwargs["step"]
    assert step.tools == tools
    assert step.ground_truth == {"answer": "expected"}


def test_step_legacy_serialization_is_unchanged_without_children():
    """A Step without children keeps the existing serialized payload."""
    step = Step(type="llm", name="chat", input="hello")

    assert step.model_dump(mode="json") == {
        "type": "llm",
        "name": "chat",
        "input": "hello",
        "output": None,
        "context": None,
        "tools": None,
        "ground_truth": None,
    }
    assert Step(type="llm", name="chat", input="hello", children=None).children is None


def test_step_accepts_and_preserves_recursive_children():
    """Nested child steps retain all provider-neutral runtime fields."""
    child_payload = {
        "type": "llm",
        "name": "child",
        "input": {"prompt": "question"},
        "output": "answer",
        "context": {"conversation": ["earlier"]},
        "tools": [{"name": "lookup", "input_schema": {"type": "object"}}],
        "ground_truth": {"answer": "expected"},
        "children": [
            {"type": "retriever", "name": "search", "input": {"query": "q"}}
        ],
    }
    parent = Step(type="llm", name="parent", input="hello", children=[child_payload])

    assert parent.children is not None
    child = parent.children[0]
    assert child.type == child_payload["type"]
    assert child.name == child_payload["name"]
    assert child.input == child_payload["input"]
    assert child.output == child_payload["output"]
    assert child.context == child_payload["context"]
    assert child.tools == child_payload["tools"]
    assert child.ground_truth == child_payload["ground_truth"]
    assert child.children is not None
    assert child.children[0].input == {"query": "q"}


def test_step_rejects_invalid_child_values_clearly():
    with pytest.raises(ValidationError, match="children"):
        Step(type="llm", name="parent", input="hello", children=["not a step"])


@pytest.mark.asyncio
async def test_evaluate_controls_passes_children_without_recursively_evaluating(monkeypatch):
    """Children are data on the current Step and are not evaluated separately."""
    mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0))
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)
    children = [
        Step(
            type="tool",
            name="search",
            input={"query": "nested"},
            children=[Step(type="llm", name="nested", input="prompt")],
        )
    ]

    with patch("agent_control.state.server_url", "http://localhost:8000"):
        await evaluation.evaluate_controls(
            step_name="parent",
            input="hello",
            children=children,
            stage="pre",
            agent_name="test-bot",
        )

    mock_check.assert_awaited_once()
    passed_step = mock_check.call_args.kwargs["step"]
    assert passed_step.children == children
    assert passed_step.children[0].children == children[0].children


@pytest.mark.asyncio
async def test_evaluate_controls_uses_session_api_key_header(monkeypatch):
    """evaluate_controls should pass init's API-key header into the client."""
    mock_result = EvaluationResult(is_safe=True, confidence=1.0)
    mock_check = AsyncMock(return_value=mock_result)
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)

    with patch("agent_control.state.server_url", "http://localhost:8000"), patch(
        "agent_control.state.api_key", "test-key"
    ), patch("agent_control.state.api_key_header", "Galileo-API-Key"):
        await evaluation.evaluate_controls(
            step_name="chat",
            input="hello",
            stage="pre",
            agent_name="test-bot",
        )

    client = mock_check.call_args.kwargs["client"]
    assert client.api_key_header == "Galileo-API-Key"


@pytest.mark.asyncio
async def test_check_evaluation_forwards_target_context():
    """When target_type and target_id are supplied, they are forwarded to the server."""

    class DummyResponse:
        def raise_for_status(self) -> None:
            return None

        def json(self) -> dict[str, object]:
            return {"is_safe": True, "confidence": 1.0}

    client = MagicMock()
    client.http_client = MagicMock()
    client.http_client.post = AsyncMock(return_value=DummyResponse())

    await evaluation.check_evaluation(
        client=client,
        agent_name="Agent-Example_01",
        step={"type": "llm", "name": "chat", "input": "hello"},
        stage="pre",
        target_type="env",
        target_id="prod",
    )

    sent = client.http_client.post.await_args.kwargs["json"]
    assert sent["target_type"] == "env"
    assert sent["target_id"] == "prod"


@pytest.mark.asyncio
async def test_evaluate_controls_forwards_target_context(monkeypatch):
    """evaluate_controls passes target_type/target_id into check_evaluation_with_local.

    Per the V1 contract, per-call target must match the session target.
    The test pins forwarding behavior with state.target_type/state.target_id
    set to the same values as the per-call args.
    """
    mock_result = EvaluationResult(is_safe=True, confidence=1.0)
    mock_check = AsyncMock(return_value=mock_result)
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)

    with patch("agent_control.state.server_url", "http://localhost:8000"), patch(
        "agent_control.state.api_key", None
    ), patch("agent_control._state.state.target_type", "env"), patch(
        "agent_control._state.state.target_id", "prod"
    ):
        await evaluation.evaluate_controls(
            step_name="chat",
            input="hello",
            stage="pre",
            agent_name="test-bot",
            target_type="env",
            target_id="prod",
        )

    kwargs = mock_check.call_args.kwargs
    assert kwargs["target_type"] == "env"
    assert kwargs["target_id"] == "prod"


@pytest.mark.asyncio
async def test_check_evaluation_does_not_default_target_from_state():
    """``check_evaluation`` is not session-bound.

    The caller supplies its own client; the helper does not consult
    ``init()``-time session state. A caller that omits target params gets
    a non-target-bearing request even when the session has a target set —
    session-target enforcement lives on the session-bound entry point
    (``evaluate_controls``).
    """

    class DummyResponse:
        def raise_for_status(self) -> None:
            return None

        def json(self) -> dict[str, object]:
            return {"is_safe": True, "confidence": 1.0}

    client = MagicMock()
    client.http_client = MagicMock()
    client.http_client.post = AsyncMock(return_value=DummyResponse())

    with patch("agent_control._state.state.target_type", "env"), patch(
        "agent_control._state.state.target_id", "prod"
    ):
        await evaluation.check_evaluation(
            client=client,
            agent_name="Agent-Example_01",
            step={"type": "llm", "name": "chat", "input": "hello"},
            stage="pre",
        )

    sent = client.http_client.post.await_args.kwargs["json"]
    assert sent["target_type"] is None
    assert sent["target_id"] is None


@pytest.mark.asyncio
async def test_check_evaluation_partial_target_pair_rejected():
    """Per-call target params must be supplied together."""
    client = MagicMock()
    client.http_client = MagicMock()
    client.http_client.post = AsyncMock()

    with pytest.raises(ValueError):
        await evaluation.check_evaluation(
            client=client,
            agent_name="Agent-Example_01",
            step={"type": "llm", "name": "chat", "input": "hello"},
            stage="pre",
            target_type="env",  # target_id missing
        )


@pytest.mark.asyncio
async def test_evaluate_controls_per_call_target_must_match_session_target(monkeypatch):
    """A per-call target that disagrees with init()'s target is rejected.

    The cached controls are fetched for the session target; accepting a
    mismatched per-call target would drive stale local-first evaluation.
    """
    mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0))
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)
    session_agent = MagicMock(name="session_agent")

    with patch("agent_control.state.server_url", "http://localhost:8000"), patch(
        "agent_control.state.api_key", None
    ), patch("agent_control._state.state.current_agent", session_agent), patch(
        "agent_control._state.state.target_type", "env"
    ), patch("agent_control._state.state.target_id", "prod"):
        with pytest.raises(ValueError, match="must match the target context fixed at init"):
            await evaluation.evaluate_controls(
                step_name="chat",
                input="hello",
                stage="pre",
                agent_name="test-bot",
                target_type="env",
                target_id="staging",  # session is "prod"
            )


@pytest.mark.asyncio
async def test_evaluate_controls_no_target_session_rejects_per_call_target(monkeypatch):
    """An init() without target context is itself a fixed (None, None) session.

    The session's cached controls were fetched without target context, so a
    per-call target that supplies an explicit (target_type, target_id) would
    drive stale local-first evaluation. Reject the mismatch.
    """
    mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0))
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)
    session_agent = MagicMock(name="session_agent")

    with patch("agent_control.state.server_url", "http://localhost:8000"), patch(
        "agent_control.state.api_key", None
    ), patch("agent_control._state.state.current_agent", session_agent), patch(
        "agent_control._state.state.target_type", None
    ), patch("agent_control._state.state.target_id", None):
        with pytest.raises(
            ValueError, match="must match the target context fixed at init"
        ):
            await evaluation.evaluate_controls(
                step_name="chat",
                input="hello",
                stage="pre",
                agent_name="test-bot",
                target_type="env",
                target_id="prod",
            )


@pytest.mark.asyncio
async def test_evaluate_controls_defaults_target_from_state(monkeypatch):
    """``evaluate_controls`` falls back to state target when params omitted."""
    mock_result = EvaluationResult(is_safe=True, confidence=1.0)
    mock_check = AsyncMock(return_value=mock_result)
    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)

    with patch("agent_control.state.server_url", "http://localhost:8000"), patch(
        "agent_control.state.api_key", None
    ), patch("agent_control._state.state.target_type", "env"), patch(
        "agent_control._state.state.target_id", "prod"
    ):
        await evaluation.evaluate_controls(
            step_name="chat",
            input="hello",
            stage="pre",
            agent_name="test-bot",
        )

    kwargs = mock_check.call_args.kwargs
    assert kwargs["target_type"] == "env"
    assert kwargs["target_id"] == "prod"