mlflow--mlflow/tests/genai/evaluate/test_utils.py

import json
import sys
from typing import Any, Literal
from unittest.mock import MagicMock, Mock, patch

import pandas as pd
import pytest

import mlflow
from mlflow.entities.assessment_source import AssessmentSource
from mlflow.entities.span import SpanType
from mlflow.entities.trace import Trace
from mlflow.exceptions import MlflowException
from mlflow.genai import scorer
from mlflow.genai.datasets import EvaluationDataset, create_dataset
from mlflow.genai.evaluation.utils import (
    _convert_scorer_to_legacy_metric,
    _convert_to_eval_set,
    _deserialize_trace_column_if_needed,
    validate_tags,
)
from mlflow.genai.scorers.builtin_scorers import RelevanceToQuery
from mlflow.utils.spark_utils import is_spark_connect_mode

from tests.genai.conftest import databricks_only


@pytest.fixture(scope="module")
def spark():
    # databricks-agents installs databricks-connect
    if is_spark_connect_mode():
        pytest.skip("Local Spark Session is not supported when databricks-connect is installed.")

    from pyspark.sql import SparkSession

    with SparkSession.builder.getOrCreate() as spark:
        yield spark


def count_rows(data: Any) -> int:
    try:
        from mlflow.utils.spark_utils import get_spark_dataframe_type

        if isinstance(data, get_spark_dataframe_type()):
            return data.count()
    except Exception:
        pass

    if isinstance(data, EvaluationDataset):
        data = data.to_df()

    return len(data)


@pytest.fixture
def sample_dict_data_single():
    return [
        {
            "inputs": {"question": "What is Spark?"},
            "outputs": "actual response for first question",
            "expectations": {"expected_response": "expected response for first question"},
            "tags": {"sample_tag": "value"},
        },
    ]


@pytest.fixture
def sample_dict_data_multiple():
    return [
        {
            "inputs": {"question": "What is Spark?"},
            "outputs": "actual response for first question",
            "expectations": {"expected_response": "expected response for first question"},
            "tags": {"category": "spark"},
        },
        {
            "inputs": {"question": "How can you minimize data shuffling in Spark?"},
            "outputs": "actual response for second question",
            "expectations": {"expected_response": "expected response for second question"},
            "tags": {"category": "spark", "topic": "optimization"},
        },
        # Some records might not have expectations or tags
        {
            "inputs": {"question": "What is MLflow?"},
            "outputs": "actual response for third question",
            "expectations": {},
            "tags": {},
        },
    ]


@pytest.fixture
def sample_dict_data_multiple_with_custom_expectations():
    return [
        {
            "inputs": {"question": "What is Spark?"},
            "outputs": "actual response for first question",
            "expectations": {
                "expected_response": "expected response for first question",
                "my_custom_expectation": "custom expectation for the first question",
            },
        },
        {
            "inputs": {"question": "How can you minimize data shuffling in Spark?"},
            "outputs": "actual response for second question",
            "expectations": {
                "expected_response": "expected response for second question",
                "my_custom_expectation": "custom expectation for the second question",
            },
        },
        # Some records might not have all expectations
        {
            "inputs": {"question": "What is MLflow?"},
            "outputs": "actual response for third question",
            "expectations": {
                "my_custom_expectation": "custom expectation for the third question",
            },
        },
    ]


@pytest.fixture
def sample_pd_data(sample_dict_data_multiple):
    """Returns a pandas DataFrame with sample data"""
    return pd.DataFrame(sample_dict_data_multiple)


@pytest.fixture
def sample_spark_data(sample_pd_data, spark):
    """Convert pandas DataFrame to PySpark DataFrame"""
    return spark.createDataFrame(sample_pd_data)


@pytest.fixture
def sample_spark_data_with_string_columns(sample_pd_data, spark):
    # Cast inputs and expectations columns to string
    df = sample_pd_data.copy()
    df["inputs"] = df["inputs"].apply(json.dumps)
    df["expectations"] = df["expectations"].apply(json.dumps)
    return spark.createDataFrame(df)


@pytest.fixture
def sample_evaluation_dataset(sample_dict_data_single):
    dataset = create_dataset("test")
    dataset.merge_records(sample_dict_data_single)
    return dataset


_ALL_DATA_FIXTURES = [
    "sample_dict_data_single",
    "sample_dict_data_multiple",
    "sample_dict_data_multiple_with_custom_expectations",
    "sample_pd_data",
    "sample_spark_data",
    "sample_spark_data_with_string_columns",
    "sample_evaluation_dataset",
]


class TestModel:
    @mlflow.trace(span_type=SpanType.AGENT)
    def predict(self, question: str) -> str:
        response = self.call_llm(messages=[{"role": "user", "content": question}])
        return response["choices"][0]["message"]["content"]

    @mlflow.trace(span_type=SpanType.LLM)
    def call_llm(self, messages: list[dict[str, Any]]) -> dict[str, Any]:
        return {"choices": [{"message": {"role": "assistant", "content": "I don't know"}}]}


def get_test_traces(type=Literal["pandas", "list"]):
    model = TestModel()

    model.predict("What is MLflow?")
    trace_id = mlflow.get_last_active_trace_id()

    # Add assessments. Since log_assessment API is not supported in OSS MLflow yet, we
    # need to add it to the trace info manually.
    source = AssessmentSource(source_id="test", source_type="HUMAN")
    # 1. Expectation with reserved name "expected_response"
    mlflow.log_expectation(
        trace_id=trace_id,
        name="expected_response",
        value="expected response for first question",
        source=source,
    )
    # 2. Expectation with reserved name "expected_facts"
    mlflow.log_expectation(
        trace_id=trace_id,
        name="expected_facts",
        value=["fact1", "fact2"],
        source=source,
    )
    # 3. Expectation with reserved name "guidelines"
    mlflow.log_expectation(
        trace_id=trace_id,
        name="guidelines",
        value=["Be polite", "Be kind"],
        source=source,
    )
    # 4. Expectation with custom name "my_custom_expectation"
    mlflow.log_expectation(
        trace_id=trace_id,
        name="my_custom_expectation",
        value="custom expectation for the first question",
        source=source,
    )
    # 5. Non-expectation assessment
    mlflow.log_feedback(
        trace_id=trace_id,
        name="feedback",
        value="some feedback",
        source=source,
    )
    traces = mlflow.search_traces(return_type=type, order_by=["timestamp_ms ASC"])
    return [{"trace": trace} for trace in traces] if type == "list" else traces


@pytest.mark.parametrize("input_type", ["list", "pandas"])
def test_convert_to_legacy_eval_traces(input_type):
    sample_data = get_test_traces(type=input_type)
    data = _convert_to_eval_set(sample_data)

    assert "trace" in data.columns

    # "inputs" column should be derived from the trace
    assert "inputs" in data.columns
    assert list(data["inputs"]) == [{"question": "What is MLflow?"}]
    assert data["expectations"][0] == {
        "expected_response": "expected response for first question",
        "expected_facts": ["fact1", "fact2"],
        "guidelines": ["Be polite", "Be kind"],
        "my_custom_expectation": "custom expectation for the first question",
    }
    # Assessment with type "Feedback" should not be present in the transformed data
    assert "feedback" not in data.columns


@pytest.mark.parametrize("data_fixture", _ALL_DATA_FIXTURES)
def test_convert_to_eval_set_has_no_errors(data_fixture, request):
    sample_data = request.getfixturevalue(data_fixture)

    transformed_data = _convert_to_eval_set(sample_data)

    assert "inputs" in transformed_data.columns
    assert "outputs" in transformed_data.columns
    assert "expectations" in transformed_data.columns


def test_convert_to_eval_set_without_request_and_response():
    for _ in range(3):
        with mlflow.start_span():
            pass

    trace_df = mlflow.search_traces()
    trace_df = trace_df[["trace"]]
    transformed_data = _convert_to_eval_set(trace_df)

    assert "inputs" in transformed_data.columns
    assert "outputs" in transformed_data.columns
    assert transformed_data["inputs"].isna().all()


def test_convert_to_eval_set_with_missing_root_span():
    # Create traces
    for _ in range(2):
        with mlflow.start_span():
            pass

    trace_df = mlflow.search_traces()
    trace_df = trace_df[["trace"]]

    # Deserialize the trace from JSON string to Trace object
    trace_df["trace"] = trace_df["trace"].apply(
        lambda t: Trace.from_json(t) if isinstance(t, str) else t
    )

    # Mock _get_root_span to return None for the first trace to simulate missing root span
    with patch.object(trace_df["trace"].iloc[0].data, "_get_root_span", return_value=None):
        transformed_data = _convert_to_eval_set(trace_df)

    # Verify inputs and outputs columns exist
    assert "inputs" in transformed_data.columns
    assert "outputs" in transformed_data.columns

    # Verify first trace has None for inputs/outputs (missing root span)
    assert transformed_data["inputs"].iloc[0] is None
    assert transformed_data["outputs"].iloc[0] is None

    # Verify second trace has None for inputs/outputs (normal empty span behavior)
    assert transformed_data["inputs"].iloc[1] is None
    assert transformed_data["outputs"].iloc[1] is None


def test_convert_to_legacy_eval_raise_for_invalid_json_columns(spark):
    # Data with invalid `inputs` column
    df = spark.createDataFrame([
        {"inputs": "invalid json", "expectations": '{"expected_response": "expected"}'},
        {"inputs": "invalid json", "expectations": '{"expected_response": "expected"}'},
    ])
    with pytest.raises(MlflowException, match="Failed to parse `inputs` column."):
        _convert_to_eval_set(df)

    # Data with invalid `expectations` column
    df = spark.createDataFrame([
        {
            "inputs": '{"question": "What is the capital of France?"}',
            "expectations": "invalid expectations",
        },
        {
            "inputs": '{"question": "What is the capital of Germany?"}',
            "expectations": "invalid expectations",
        },
    ])
    with pytest.raises(MlflowException, match="Failed to parse `expectations` column."):
        _convert_to_eval_set(df)


def _trace_test_cases():
    data = {
        "info": {
            "trace_id": "test-trace-id",
            "trace_location": {
                "type": "MLFLOW_EXPERIMENT",
                "mlflow_experiment": {"experiment_id": "0"},
            },
            "request_time": "2024-01-21T12:00:00Z",
            "state": "OK",
            "trace_metadata": {},
            "tags": {},
            "assessments": [],
        },
        "data": {"spans": []},
    }
    return [
        pytest.param(data, dict, id="dict"),
        pytest.param(json.dumps(data), str, id="string"),
        pytest.param(Trace.from_dict(data), Trace, id="trace_object"),
    ]


@pytest.mark.parametrize(("trace_value", "expected_input_type"), _trace_test_cases())
def test_deserialize_trace_column(trace_value, expected_input_type):
    df = pd.DataFrame([{"trace": trace_value, "inputs": {"question": "test"}}])
    assert isinstance(df["trace"].iloc[0], expected_input_type)

    result = _deserialize_trace_column_if_needed(df)
    assert isinstance(result["trace"].iloc[0], Trace)
    assert result["trace"].iloc[0].info.trace_id == "test-trace-id"


def test_deserialize_trace_column_with_none():
    df = pd.DataFrame([{"trace": None, "inputs": {"question": "test"}}])

    result = _deserialize_trace_column_if_needed(df)
    assert result["trace"].iloc[0] is None


@pytest.mark.parametrize("data_fixture", _ALL_DATA_FIXTURES)
def test_scorer_receives_correct_data(data_fixture, request):
    sample_data = request.getfixturevalue(data_fixture)

    received_args = []

    @scorer
    def dummy_scorer(inputs, outputs, expectations):
        received_args.append((
            inputs["question"],
            outputs,
            expectations.get("expected_response"),
            expectations.get("my_custom_expectation"),
        ))
        return 0

    mlflow.genai.evaluate(
        data=sample_data,
        scorers=[dummy_scorer],
    )

    all_inputs, all_outputs, all_expectations, all_custom_expectations = zip(*received_args)
    row_count = count_rows(sample_data)
    expected_inputs = [
        "What is Spark?",
        "How can you minimize data shuffling in Spark?",
        "What is MLflow?",
    ][:row_count]
    expected_outputs = [
        "actual response for first question",
        "actual response for second question",
        "actual response for third question",
    ][:row_count]
    expected_expectations = [
        "expected response for first question",
        "expected response for second question",
        None,
    ][:row_count]

    assert set(all_inputs) == set(expected_inputs)
    assert set(all_outputs) == set(expected_outputs)
    assert set(all_expectations) == set(expected_expectations)

    if data_fixture == "sample_dict_data_multiple_with_custom_expectations":
        expected_custom_expectations = [
            "custom expectation for the first question",
            "custom expectation for the second question",
            "custom expectation for the third question",
        ]
        assert set(all_custom_expectations) == set(expected_custom_expectations)


def test_input_is_required_if_trace_is_not_provided():
    with patch("mlflow.genai.evaluation.harness.run") as mock_evaluate:
        with pytest.raises(MlflowException, match="inputs.*required"):
            mlflow.genai.evaluate(
                data=pd.DataFrame({"outputs": ["Paris"]}),
                scorers=[RelevanceToQuery()],
            )

        mock_evaluate.assert_not_called()

        mlflow.genai.evaluate(
            data=pd.DataFrame({
                "inputs": [{"question": "What is the capital of France?"}],
                "outputs": ["Paris"],
            }),
            scorers=[RelevanceToQuery()],
        )
        mock_evaluate.assert_called_once()


def test_input_is_optional_if_trace_is_provided():
    with mlflow.start_span() as span:
        span.set_inputs({"question": "What is the capital of France?"})
        span.set_outputs("Paris")

    trace = mlflow.get_trace(span.trace_id)

    with patch("mlflow.genai.evaluation.harness.run") as mock_evaluate:
        mlflow.genai.evaluate(
            data=pd.DataFrame({"trace": [trace]}),
            scorers=[RelevanceToQuery()],
        )

        mock_evaluate.assert_called_once()


@pytest.mark.parametrize("input_type", ["list", "pandas"])
def test_scorer_receives_correct_data_with_trace_data(input_type, monkeypatch: pytest.MonkeyPatch):
    sample_data = get_test_traces(type=input_type)
    received_args = []

    @scorer
    def dummy_scorer(inputs, outputs, expectations, trace):
        received_args.append((inputs, outputs, expectations, trace))
        return 0

    # Disable logging traces to MLflow to avoid calling mlflow APIs which need to be mocked
    monkeypatch.setenv("AGENT_EVAL_LOG_TRACES_TO_MLFLOW_ENABLED", "false")
    mlflow.genai.evaluate(
        data=sample_data,
        scorers=[dummy_scorer],
    )

    inputs, outputs, expectations, trace = received_args[0]
    assert inputs == {"question": "What is MLflow?"}
    assert outputs == "I don't know"
    assert expectations == {
        "expected_response": "expected response for first question",
        "expected_facts": ["fact1", "fact2"],
        "guidelines": ["Be polite", "Be kind"],
        "my_custom_expectation": "custom expectation for the first question",
    }
    assert isinstance(trace, Trace)


@pytest.mark.parametrize("data_fixture", _ALL_DATA_FIXTURES)
def test_predict_fn_receives_correct_data(data_fixture, request):
    sample_data = request.getfixturevalue(data_fixture)

    received_args = []

    def predict_fn(question: str):
        received_args.append(question)
        return question

    @scorer
    def dummy_scorer(inputs, outputs):
        return 0

    mlflow.genai.evaluate(
        predict_fn=predict_fn,
        data=sample_data,
        scorers=[dummy_scorer],
    )

    received_args.pop(0)  # Remove the one-time prediction to check if a model is traced
    row_count = count_rows(sample_data)
    assert len(received_args) == row_count
    expected_contents = [
        "What is Spark?",
        "How can you minimize data shuffling in Spark?",
        "What is MLflow?",
    ][:row_count]
    # Using set because eval harness runs predict_fn in parallel
    assert set(received_args) == set(expected_contents)


def test_convert_scorer_to_legacy_metric_aggregations_attribute(monkeypatch):
    mock_metric_instance = MagicMock()

    # NB: Mocking the behavior of databricks-agents, which does not have the aggregations
    # argument for the evaluation interface for a metric.
    def mock_metric_decorator(**kwargs):
        if "aggregations" in kwargs:
            raise TypeError("metric() got an unexpected keyword argument 'aggregations'")
        assert set(kwargs.keys()) <= {"eval_fn", "name"}
        return mock_metric_instance

    mock_evals = Mock(metric=mock_metric_decorator)
    mock_evals.judges = Mock()  # Add the judges submodule to prevent AttributeError

    monkeypatch.setitem(sys.modules, "databricks.agents.evals", mock_evals)
    monkeypatch.setitem(sys.modules, "databricks.agents.evals.judges", mock_evals.judges)

    mock_scorer = Mock()
    mock_scorer.name = "test_scorer"
    mock_scorer.aggregations = ["mean", "max", "p90"]
    mock_scorer.run = Mock(return_value={"score": 1.0})

    result = _convert_scorer_to_legacy_metric(mock_scorer)

    assert result.aggregations == ["mean", "max", "p90"]


@databricks_only
def test_convert_scorer_to_legacy_metric():
    # Test with a built-in scorer
    builtin_scorer = RelevanceToQuery()
    legacy_metric = _convert_scorer_to_legacy_metric(builtin_scorer)

    # Verify the metric has the _is_builtin_scorer attribute set to True
    assert hasattr(legacy_metric, "_is_builtin_scorer")
    assert legacy_metric._is_builtin_scorer is True
    assert legacy_metric.name == builtin_scorer.name

    # Test with a custom scorer
    @scorer(name="custom_scorer", aggregations=["mean", "max"])
    def custom_scorer_func(inputs, outputs=None, expectations=None, **kwargs):
        return {"score": 1.0}

    custom_scorer_instance = custom_scorer_func
    legacy_metric_custom = _convert_scorer_to_legacy_metric(custom_scorer_instance)

    # Verify the metric has the _is_builtin_scorer attribute set to False
    assert hasattr(legacy_metric_custom, "_is_builtin_scorer")
    assert legacy_metric_custom._is_builtin_scorer is False
    assert legacy_metric_custom.name == custom_scorer_instance.name
    assert legacy_metric_custom.aggregations == custom_scorer_instance.aggregations


@pytest.mark.parametrize(
    "aggregations",
    [
        ["mean", "max", "mean", "median", "variance", "p90"],
        [sum, max],
    ],
)
@databricks_only
def test_scorer_pass_through_aggregations(aggregations):
    @scorer(name="custom_scorer", aggregations=aggregations)
    def custom_scorer_func(outputs):
        return {"score": 1.0}

    legacy_metric_custom = _convert_scorer_to_legacy_metric(custom_scorer_func)
    assert legacy_metric_custom.name == "custom_scorer"
    assert legacy_metric_custom.aggregations == aggregations

    builtin_scorer = RelevanceToQuery(aggregations=aggregations)
    legacy_metric_builtin = _convert_scorer_to_legacy_metric(builtin_scorer)
    assert legacy_metric_builtin.name == "relevance_to_query"
    assert legacy_metric_builtin.aggregations == builtin_scorer.aggregations


@pytest.mark.parametrize(
    "tags",
    [
        None,
        {},
        {"key": "value"},
        {"env": "test", "model": "v1.0"},
        {"key": 123},  # Values can be any type
        {"key1": "value1", "key2": None},  # Values can be any type
    ],
)
def test_validate_tags_valid(tags):
    validate_tags(tags)


@pytest.mark.parametrize(
    ("tags", "expected_error"),
    [
        ("invalid", "Tags must be a dictionary, got str"),
        (123, "Tags must be a dictionary, got int"),
        ([1, 2, 3], "Tags must be a dictionary, got list"),
        ({123: "value"}, "Invalid tags:\n  - Key 123 has type int; expected str."),
        (
            {"key1": "value1", 123: "value2"},
            "Invalid tags:\n  - Key 123 has type int; expected str.",
        ),
        (
            {123: "value1", 456: "value2"},
            (
                "Invalid tags:\n  - Key 123 has type int; expected str."
                "\n  - Key 456 has type int; expected str."
            ),
        ),
    ],
)
def test_validate_tags_invalid(tags, expected_error):
    with pytest.raises(MlflowException, match=expected_error):
        validate_tags(tags)