/tests/genai/optimize/test_util.py

from typing import Any, Union

import pytest
from pydantic import BaseModel

from mlflow.entities.assessment import Feedback
from mlflow.exceptions import MlflowException
from mlflow.genai.judges import CategoricalRating
from mlflow.genai.optimize.util import (
    create_metric_from_scorers,
    infer_type_from_value,
    validate_train_data,
)
from mlflow.genai.scorers import scorer


@pytest.mark.parametrize(
    ("input_value", "expected_type"),
    [
        (None, type(None)),
        (True, bool),
        (42, int),
        (3.14, float),
        ("hello", str),
    ],
)
def test_infer_primitive_types(input_value, expected_type):
    assert infer_type_from_value(input_value) == expected_type


@pytest.mark.parametrize(
    ("input_list", "expected_type"),
    [
        ([], list[Any]),
        ([1, 2, 3], list[int]),
        (["a", "b", "c"], list[str]),
        ([1.0, 2.0, 3.0], list[float]),
        ([True, False, True], list[bool]),
        ([1, "hello", True], list[Union[int, str, bool]]),  # noqa: UP007
        ([1, "hello", True], list[int | str | bool]),
        ([1, 2.0], list[int | float]),
        ([[1, 2], [3, 4]], list[list[int]]),
        ([["a"], ["b", "c"]], list[list[str]]),
    ],
)
def test_infer_list_types(input_list, expected_type):
    assert infer_type_from_value(input_list) == expected_type


@pytest.mark.parametrize(
    ("input_dict", "expected_fields"),
    [
        ({"name": "John", "age": 30, "active": True}, {"name": str, "age": int, "active": bool}),
        ({"score": 95.5, "passed": True}, {"score": float, "passed": bool}),
    ],
)
def test_infer_simple_dict(input_dict, expected_fields):
    result = infer_type_from_value(input_dict)

    assert isinstance(result, type)
    assert issubclass(result, BaseModel)

    for field_name, expected_type in expected_fields.items():
        assert result.__annotations__[field_name] == expected_type


def test_infer_nested_dict():
    data = {
        "user": {"name": "John", "scores": [85, 90, 95]},
        "settings": {"enabled": True, "theme": "dark"},
    }
    result = infer_type_from_value(data)

    assert isinstance(result, type)
    assert issubclass(result, BaseModel)

    # Check nested model types
    user_model = result.__annotations__["user"]
    settings_model = result.__annotations__["settings"]

    assert issubclass(user_model, BaseModel)
    assert issubclass(settings_model, BaseModel)

    # Check nested field types
    assert user_model.__annotations__["name"] == str
    assert user_model.__annotations__["scores"] == list[int]
    assert settings_model.__annotations__["enabled"] == bool
    assert settings_model.__annotations__["theme"] == str


@pytest.mark.parametrize(
    ("model_class", "model_data"),
    [
        (
            type("UserModel", (BaseModel,), {"__annotations__": {"name": str, "age": int}}),
            {"name": "John", "age": 30},
        ),
        (
            type("ProductModel", (BaseModel,), {"__annotations__": {"id": int, "price": float}}),
            {"id": 1, "price": 99.99},
        ),
    ],
)
def test_infer_pydantic_model(model_class, model_data):
    model = model_class(**model_data)
    result = infer_type_from_value(model)
    assert result == model_class


@pytest.mark.parametrize(
    "type_to_infer",
    [
        type("CustomClass", (), {}),
        type("AnotherClass", (), {"custom_attr": 42}),
    ],
)
def test_infer_unsupported_type(type_to_infer):
    obj = type_to_infer()
    assert infer_type_from_value(obj) == Any


@pytest.mark.parametrize(
    ("input_dict", "model_name"),
    [
        ({"name": "John", "age": 30}, "UserData"),
        ({"id": 1, "value": "test"}, "TestModel"),
    ],
)
def test_model_name_parameter(input_dict, model_name):
    result = infer_type_from_value(input_dict, model_name=model_name)
    assert result.__name__ == model_name


@pytest.mark.parametrize(
    ("score", "expected_score"),
    [
        (CategoricalRating.YES, 1.0),
        (CategoricalRating.NO, 0.0),
        ("yes", 1.0),
        ("no", 0.0),
        (True, 1.0),
        (False, 0.0),
        (1, 1.0),
        (0, 0.0),
        (1.0, 1.0),
        (0.0, 0.0),
    ],
)
def test_create_metric_from_scorers_with_single_score(score, expected_score):
    @scorer(name="test_scorer")
    def test_scorer(inputs, outputs):
        return Feedback(name="test_scorer", value=score, rationale="test rationale")

    metric = create_metric_from_scorers([test_scorer])

    result = metric({"input": "test"}, {"output": "result"}, {}, None)
    assert result[0] == expected_score
    assert result[1] == {"test_scorer": "test rationale"}
    assert result[2] == {"test_scorer": expected_score}


def test_create_metric_from_scorers_with_multiple_categorical_ratings():
    @scorer(name="scorer1")
    def scorer1(inputs, outputs):
        return Feedback(name="scorer1", value=CategoricalRating.YES, rationale="rationale1")

    @scorer(name="scorer2")
    def scorer2(inputs, outputs):
        return Feedback(name="scorer2", value=CategoricalRating.YES, rationale="rationale2")

    metric = create_metric_from_scorers([scorer1, scorer2])

    # Should average: (1.0 + 1.0) / 2 = 1.0
    result = metric({"input": "test"}, {"output": "result"}, {}, None)
    assert result[0] == 1.0
    assert result[1] == {"scorer1": "rationale1", "scorer2": "rationale2"}
    assert result[2] == {"scorer1": 1.0, "scorer2": 1.0}


@pytest.mark.parametrize(
    ("train_data", "scorers", "expected_error"),
    [
        # Empty inputs
        (
            [{"inputs": {}, "outputs": "result"}],
            [],
            "Record 0 is missing required 'inputs' field or it is empty",
        ),
        # Missing inputs
        (
            [{"outputs": "result"}],
            [],
            "Record 0 is missing required 'inputs' field or it is empty",
        ),
    ],
)
def test_validate_train_data_errors(train_data, scorers, expected_error):
    import pandas as pd

    with pytest.raises(MlflowException, match=expected_error):
        validate_train_data(pd.DataFrame(train_data), scorers, lambda **kwargs: None)


@pytest.mark.parametrize(
    "train_data",
    [
        # Valid with outputs
        [{"inputs": {"text": "hello"}, "outputs": "result"}],
        # Valid with expectations
        [{"inputs": {"text": "hello"}, "expectations": {"expected": "result"}}],
        # Multiple valid records
        [
            {"inputs": {"text": "hello"}, "outputs": "result1"},
            {"inputs": {"text": "world"}, "expectations": {"expected": "result2"}},
        ],
        # Falsy but valid values: False as output
        [{"inputs": {"text": "hello"}, "outputs": False}],
    ],
)
def test_validate_train_data_success(train_data):
    import pandas as pd

    validate_train_data(pd.DataFrame(train_data), [], lambda **kwargs: None)