mlflow--mlflow/tests/genai/optimize/test_job.py

import sys
from unittest import mock

import pytest

import mlflow
from mlflow.exceptions import MlflowException
from mlflow.genai.optimize.job import (
    OptimizerType,
    _build_predict_fn,
    _create_optimizer,
    _load_scorers,
    optimize_prompts_job,
)
from mlflow.genai.optimize.optimizers import GepaPromptOptimizer, MetaPromptOptimizer
from mlflow.genai.scorers import scorer
from mlflow.genai.scorers.builtin_scorers import Correctness, Safety
from mlflow.protos.prompt_optimization_pb2 import (
    OPTIMIZER_TYPE_GEPA,
    OPTIMIZER_TYPE_METAPROMPT,
    OPTIMIZER_TYPE_UNSPECIFIED,
)


def test_create_gepa_optimizer_success():
    config = {"reflection_model": "openai:/gpt-4o", "max_metric_calls": 50}
    optimizer = _create_optimizer("gepa", config)
    assert isinstance(optimizer, GepaPromptOptimizer)
    assert optimizer.reflection_model == "openai:/gpt-4o"
    assert optimizer.max_metric_calls == 50


def test_create_gepa_optimizer_case_insensitive():
    config = {"reflection_model": "openai:/gpt-4o"}
    optimizer = _create_optimizer("GEPA", config)
    assert isinstance(optimizer, GepaPromptOptimizer)


def test_create_gepa_optimizer_missing_reflection_model():
    config = {"max_metric_calls": 50}
    with pytest.raises(MlflowException, match="'reflection_model' must be specified"):
        _create_optimizer("gepa", config)


def test_create_metaprompt_optimizer_success():
    config = {"reflection_model": "openai:/gpt-4o", "guidelines": "Be concise"}
    optimizer = _create_optimizer("metaprompt", config)
    assert isinstance(optimizer, MetaPromptOptimizer)


def test_create_metaprompt_optimizer_missing_reflection_model():
    config = {"guidelines": "Be concise"}
    with pytest.raises(MlflowException, match="'reflection_model' must be specified"):
        _create_optimizer("metaprompt", config)


def test_create_optimizer_unsupported_type():
    with pytest.raises(MlflowException, match="Unsupported optimizer type: 'invalid'"):
        _create_optimizer("invalid", None)


@pytest.mark.parametrize(
    ("proto_value", "expected_type", "expected_str", "error_match"),
    [
        (OPTIMIZER_TYPE_GEPA, OptimizerType.GEPA, "gepa", None),
        (OPTIMIZER_TYPE_METAPROMPT, OptimizerType.METAPROMPT, "metaprompt", None),
        (OPTIMIZER_TYPE_UNSPECIFIED, None, None, "optimizer_type is required"),
        (999, None, None, "Unsupported optimizer_type value"),
    ],
)
def test_optimizer_type_from_proto(proto_value, expected_type, expected_str, error_match):
    if error_match:
        with pytest.raises(MlflowException, match=error_match):
            OptimizerType.from_proto(proto_value)
    else:
        result = OptimizerType.from_proto(proto_value)
        assert result == expected_type
        assert result == expected_str


@pytest.mark.parametrize(
    ("optimizer_type", "expected_proto"),
    [
        (OptimizerType.GEPA, OPTIMIZER_TYPE_GEPA),
        (OptimizerType.METAPROMPT, OPTIMIZER_TYPE_METAPROMPT),
    ],
)
def test_optimizer_type_to_proto(optimizer_type, expected_proto):
    assert optimizer_type.to_proto() == expected_proto


def test_load_builtin_scorers():
    scorers = _load_scorers(["Correctness", "Safety"], "exp-123")

    assert len(scorers) == 2
    assert isinstance(scorers[0], Correctness)
    assert isinstance(scorers[1], Safety)


def test_load_custom_scorers():
    with (
        mock.patch("mlflow.genai.scorers.base.is_databricks_uri", return_value=True),
    ):
        experiment_id = mlflow.create_experiment("test_load_custom_scorers")

        @scorer
        def custom_scorer_1(outputs) -> bool:
            return len(outputs) > 0

        @scorer
        def custom_scorer_2(outputs) -> bool:
            return len(outputs) > 0

        custom_scorer_1.register(experiment_id=experiment_id, name="custom_scorer_1")
        custom_scorer_2.register(experiment_id=experiment_id, name="custom_scorer_2")

        scorers = _load_scorers(["custom_scorer_1", "custom_scorer_2"], experiment_id)

        assert len(scorers) == 2
        assert scorers[0].name == "custom_scorer_1"
        assert scorers[1].name == "custom_scorer_2"

        mlflow.delete_experiment(experiment_id)


def test_load_scorer_not_found_raises_error():
    experiment_id = mlflow.create_experiment("test_load_scorer_not_found")

    with pytest.raises(MlflowException, match="Scorer 'unknown_scorer' not found"):
        _load_scorers(["unknown_scorer"], experiment_id)

    mlflow.delete_experiment(experiment_id)


def test_build_predict_fn_success():
    mock_prompt = mock.MagicMock()
    mock_prompt.model_config = {"provider": "openai", "model_name": "gpt-4o"}
    mock_prompt.format.return_value = "formatted prompt"

    mock_litellm = mock.MagicMock()
    mock_response = mock.MagicMock()
    mock_response.choices = [mock.MagicMock()]
    mock_response.choices[0].message.content = "response text"
    mock_litellm.completion.return_value = mock_response

    with (
        mock.patch("mlflow.genai.optimize.job.load_prompt", return_value=mock_prompt),
        mock.patch.dict("sys.modules", {"litellm": mock_litellm}),
    ):
        predict_fn = _build_predict_fn("prompts:/test/1")
        result = predict_fn(question="What is AI?")

        assert result == "response text"
        mock_litellm.completion.assert_called_once()
        call_args = mock_litellm.completion.call_args
        assert call_args.kwargs["model"] == "openai/gpt-4o"
        mock_prompt.format.assert_called_with(question="What is AI?")


def test_build_predict_fn_missing_model_config():
    mock_prompt = mock.MagicMock()
    mock_prompt.model_config = None

    mock_litellm = mock.MagicMock()

    with (
        mock.patch("mlflow.genai.optimize.job.load_prompt", return_value=mock_prompt),
        mock.patch.dict("sys.modules", {"litellm": mock_litellm}),
    ):
        with pytest.raises(MlflowException, match="doesn't have a model configuration"):
            _build_predict_fn("prompts:/test/1")


def test_build_predict_fn_missing_provider():
    mock_prompt = mock.MagicMock()
    mock_prompt.model_config = {"model_name": "gpt-4o"}

    mock_litellm = mock.MagicMock()

    with (
        mock.patch("mlflow.genai.optimize.job.load_prompt", return_value=mock_prompt),
        mock.patch.dict("sys.modules", {"litellm": mock_litellm}),
    ):
        with pytest.raises(MlflowException, match="doesn't have a model configuration"):
            _build_predict_fn("prompts:/test/1")


def test_build_predict_fn_missing_litellm():
    # Simulate litellm not being installed
    with mock.patch.dict(sys.modules, {"litellm": None}):
        with pytest.raises(
            MlflowException, match="'litellm' package is required for prompt optimization"
        ):
            _build_predict_fn("prompts:/test/1")


def test_optimize_prompts_job_has_metadata():
    assert hasattr(optimize_prompts_job, "_job_fn_metadata")
    metadata = optimize_prompts_job._job_fn_metadata
    assert metadata.name == "optimize_prompts"
    assert metadata.max_workers == 2


def test_optimize_prompts_job_calls():
    mock_dataset = mock.MagicMock()

    mock_prompt = mock.MagicMock()
    mock_prompt.model_config = {"provider": "openai", "model_name": "gpt-4o"}

    mock_optimizer = mock.MagicMock()
    mock_loaded_scorers = [mock.MagicMock(), mock.MagicMock()]
    mock_predict_fn = mock.MagicMock()

    mock_result = mock.MagicMock()
    mock_result.optimized_prompts = [mock.MagicMock()]
    mock_result.optimized_prompts[0].uri = "prompts:/test/2"
    mock_result.optimizer_name = "GepaPromptOptimizer"
    mock_result.initial_eval_score = 0.5
    mock_result.final_eval_score = 0.9

    optimizer_config = {"reflection_model": "openai:/gpt-4o"}

    with (
        mock.patch("mlflow.genai.optimize.job.get_dataset", return_value=mock_dataset),
        mock.patch("mlflow.genai.optimize.job.load_prompt", return_value=mock_prompt),
        mock.patch(
            "mlflow.genai.optimize.job._create_optimizer", return_value=mock_optimizer
        ) as mock_create_optimizer,
        mock.patch(
            "mlflow.genai.optimize.job._load_scorers", return_value=mock_loaded_scorers
        ) as mock_load_scorers,
        mock.patch(
            "mlflow.genai.optimize.job._build_predict_fn", return_value=mock_predict_fn
        ) as mock_build_predict_fn,
        mock.patch("mlflow.genai.optimize.job.set_experiment"),
        mock.patch("mlflow.genai.optimize.job.start_run"),
        mock.patch("mlflow.genai.optimize.job.MlflowClient"),
        mock.patch(
            "mlflow.genai.optimize.job.optimize_prompts", return_value=mock_result
        ) as mock_optimize_prompts,
    ):
        optimize_prompts_job(
            run_id="run-123",
            experiment_id="exp-123",
            prompt_uri="prompts:/test/1",
            dataset_id="dataset-123",
            optimizer_type="gepa",
            optimizer_config=optimizer_config,
            scorer_names=["Correctness", "Safety"],
        )

        # Verify _create_optimizer was called with correct args
        mock_create_optimizer.assert_called_once_with("gepa", optimizer_config)

        # Verify _load_scorers was called with correct args
        mock_load_scorers.assert_called_once_with(["Correctness", "Safety"], "exp-123")

        # Verify _build_predict_fn was called with correct args
        mock_build_predict_fn.assert_called_once_with("prompts:/test/1")

        # Verify optimize_prompts was called with correct args
        mock_optimize_prompts.assert_called_once_with(
            predict_fn=mock_predict_fn,
            train_data=mock_dataset,
            prompt_uris=["prompts:/test/1"],
            optimizer=mock_optimizer,
            scorers=mock_loaded_scorers,
            enable_tracking=True,
        )


def test_optimize_prompts_job_result_structure():
    mock_dataset = mock.MagicMock()

    mock_prompt = mock.MagicMock()
    mock_prompt.model_config = {"provider": "openai", "model_name": "gpt-4o"}

    mock_optimizer = mock.MagicMock()
    mock_result = mock.MagicMock()
    mock_result.optimized_prompts = [mock.MagicMock()]
    mock_result.optimized_prompts[0].uri = "prompts:/test/2"
    mock_result.optimizer_name = "GepaPromptOptimizer"
    mock_result.initial_eval_score = 0.5
    mock_result.final_eval_score = 0.9

    optimizer_config = {"reflection_model": "openai:/gpt-4o"}

    with (
        mock.patch("mlflow.genai.optimize.job.get_dataset", return_value=mock_dataset),
        mock.patch("mlflow.genai.optimize.job.load_prompt", return_value=mock_prompt),
        mock.patch("mlflow.genai.optimize.job._create_optimizer", return_value=mock_optimizer),
        mock.patch("mlflow.genai.optimize.job._load_scorers", return_value=[mock.MagicMock()]),
        mock.patch("mlflow.genai.optimize.job._build_predict_fn", return_value=lambda **k: "r"),
        mock.patch("mlflow.genai.optimize.job.set_experiment"),
        mock.patch("mlflow.genai.optimize.job.start_run"),
        mock.patch("mlflow.genai.optimize.job.MlflowClient"),
        mock.patch("mlflow.genai.optimize.job.optimize_prompts", return_value=mock_result),
    ):
        result = optimize_prompts_job(
            run_id="run-123",
            experiment_id="exp-123",
            prompt_uri="prompts:/test/1",
            dataset_id="dataset-123",
            optimizer_type="gepa",
            optimizer_config=optimizer_config,
            scorer_names=["Correctness", "Safety"],
        )

        # Verify result structure (returned as dict from asdict())
        assert result["run_id"] == "run-123"
        assert result["source_prompt_uri"] == "prompts:/test/1"
        assert result["optimized_prompt_uri"] == "prompts:/test/2"
        assert result["optimizer_name"] == "GepaPromptOptimizer"
        assert result["initial_eval_score"] == 0.5
        assert result["final_eval_score"] == 0.9
        assert result["dataset_id"] == "dataset-123"
        assert result["scorer_names"] == ["Correctness", "Safety"]