/tests/genai/judges/test_alignment_optimizer.py

from unittest.mock import Mock, patch

import pytest

from mlflow.entities.trace import Trace
from mlflow.genai.judges import AlignmentOptimizer, Judge, make_judge
from mlflow.genai.judges.base import JudgeField
from mlflow.genai.judges.optimizers import MemAlignOptimizer
from mlflow.genai.judges.utils import get_default_optimizer
from mlflow.genai.scorers import UserFrustration


class MockJudge(Judge):
    """Mock Judge implementation for testing."""

    def __init__(self, name: str = "mock_judge", **kwargs):
        super().__init__(name=name, **kwargs)

    @property
    def instructions(self) -> str:
        return f"Mock judge implementation: {self.name}"

    @property
    def feedback_value_type(self):
        return bool

    def get_input_fields(self) -> list[JudgeField]:
        """Get input fields for mock judge."""
        return [
            JudgeField(name="input", description="Mock input field"),
            JudgeField(name="output", description="Mock output field"),
        ]

    def __call__(self, **kwargs):
        from mlflow.entities.assessment import Feedback

        return Feedback(name=self.name, value=True, rationale="Mock evaluation")


class MockOptimizer(AlignmentOptimizer):
    """Mock AlignmentOptimizer implementation for testing."""

    def align(self, judge: Judge, traces: list[Trace]) -> Judge:
        # Return a new judge with modified name to show it was processed
        return MockJudge(name=f"{judge.name}_optimized")


def test_alignment_optimizer_abstract():
    with pytest.raises(TypeError, match="Can't instantiate abstract class AlignmentOptimizer"):
        AlignmentOptimizer()


def test_alignment_optimizer_align_method_required():
    class IncompleteOptimizer(AlignmentOptimizer):
        pass

    with pytest.raises(TypeError, match="Can't instantiate abstract class IncompleteOptimizer"):
        IncompleteOptimizer()


def test_concrete_optimizer_implementation():
    optimizer = MockOptimizer()
    judge = MockJudge(name="test_judge")
    traces = []  # Empty traces for testing

    # Should not raise any errors
    result = optimizer.align(judge, traces)

    assert isinstance(result, Judge)
    assert result.name == "test_judge_optimized"


class MockOptimizerWithTracking(AlignmentOptimizer):
    """Mock AlignmentOptimizer implementation with call tracking for integration tests."""

    def __init__(self):
        self.align_called = False
        self.align_args = None

    def align(self, judge: Judge, traces: list[Trace]) -> Judge:
        self.align_called = True
        self.align_args = (judge, traces)

        # Return a new judge with modified name to show it was processed
        return MockJudge(name=f"{judge.name}_aligned")


def create_mock_traces():
    """Create mock traces for testing."""
    # Create minimal mock traces - just enough to pass type checking
    mock_trace = Mock(spec=Trace)
    return [mock_trace]


def test_judge_align_method():
    judge = MockJudge(name="test_judge")
    optimizer = MockOptimizerWithTracking()
    # Replace the align method with a Mock to use built-in mechanisms
    optimizer.align = Mock(return_value=MockJudge(name="test_judge_aligned"))
    traces = create_mock_traces()

    optimized = judge.align(traces, optimizer=optimizer)

    # Verify the result
    assert isinstance(optimized, Judge)
    assert optimized.name == "test_judge_aligned"

    # Assert that optimizer.align was called with correct parameters using Mock's mechanisms
    optimizer.align.assert_called_once_with(judge, traces)


def test_judge_align_method_delegation():
    judge = MockJudge()

    # Create a spy optimizer that records calls
    optimizer = Mock(spec=AlignmentOptimizer)
    expected_result = MockJudge(name="expected")
    optimizer.align.return_value = expected_result

    traces = create_mock_traces()

    result = judge.align(traces, optimizer=optimizer)

    # Verify delegation
    optimizer.align.assert_called_once_with(judge, traces)
    assert result is expected_result


def test_judge_align_with_default_optimizer():
    judge = MockJudge()
    traces = create_mock_traces()

    # Mock the get_default_optimizer function to return our mock
    expected_result = MockJudge(name="aligned_with_default")
    mock_optimizer = Mock(spec=AlignmentOptimizer)
    mock_optimizer.align.return_value = expected_result

    with patch("mlflow.genai.judges.base.get_default_optimizer", return_value=mock_optimizer):
        result = judge.align(traces)

    # Verify delegation to default optimizer
    mock_optimizer.align.assert_called_once_with(judge, traces)
    assert result is expected_result


def test_get_default_optimizer_returns_memalign():
    optimizer = get_default_optimizer()
    reference = MemAlignOptimizer()
    assert isinstance(optimizer, MemAlignOptimizer)
    assert optimizer._retrieval_k == reference._retrieval_k
    assert optimizer._embedding_dim == reference._embedding_dim
    assert optimizer._reflection_lm == reference._reflection_lm
    assert optimizer._embedding_model == reference._embedding_model


def test_judge_align_uses_memalign_by_default():
    judge = MockJudge()
    traces = create_mock_traces()
    expected_result = MockJudge(name="aligned_with_memalign")

    with patch.object(MemAlignOptimizer, "align", return_value=expected_result) as mock_align:
        result = judge.align(traces)

    mock_align.assert_called_once_with(judge, traces)
    assert result is expected_result


def test_session_level_scorer_alignment_raises_error():
    traces = []

    conversation_judge = make_judge(
        name="conversation_judge",
        instructions="Evaluate if the {{ conversation }} is productive",
        model="openai:/gpt-4",
    )
    assert conversation_judge.is_session_level_scorer is True

    with pytest.raises(
        NotImplementedError, match="Alignment is not supported for session-level scorers"
    ):
        conversation_judge.align(traces)

    user_frustration_scorer = UserFrustration()
    assert user_frustration_scorer.is_session_level_scorer is True

    with pytest.raises(
        NotImplementedError, match="Alignment is not supported for session-level scorers"
    ):
        user_frustration_scorer.align(traces)