/tests/genai/datasets/test_fluent.py

import json
import os
import sys
import warnings
from unittest import mock

import pandas as pd
import pytest

import mlflow
from mlflow.data import Dataset
from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin
from mlflow.entities.dataset_record_source import DatasetRecordSourceType
from mlflow.entities.evaluation_dataset import (
    EvaluationDataset as EntityEvaluationDataset,
)
from mlflow.exceptions import MlflowException
from mlflow.genai.datasets import (
    EvaluationDataset,
    create_dataset,
    delete_dataset,
    delete_dataset_tag,
    get_dataset,
    search_datasets,
    set_dataset_tags,
)
from mlflow.genai.datasets.evaluation_dataset import (
    EvaluationDataset as WrapperEvaluationDataset,
)
from mlflow.store.entities.paged_list import PagedList
from mlflow.store.tracking import SEARCH_EVALUATION_DATASETS_MAX_RESULTS
from mlflow.tracking import MlflowClient
from mlflow.utils.mlflow_tags import MLFLOW_USER


@pytest.fixture
def mock_client():
    with (
        mock.patch("mlflow.tracking.client.MlflowClient") as mock_client_class,
        mock.patch("mlflow.genai.datasets.MlflowClient", mock_client_class),
    ):
        mock_client_instance = mock_client_class.return_value
        yield mock_client_instance


@pytest.fixture
def mock_databricks_environment():
    with mock.patch("mlflow.genai.datasets.is_databricks_uri", return_value=True):
        yield


@pytest.fixture
def client(db_uri):
    original_tracking_uri = mlflow.get_tracking_uri()
    mlflow.set_tracking_uri(db_uri)
    yield MlflowClient(tracking_uri=db_uri)
    mlflow.set_tracking_uri(original_tracking_uri)


@pytest.fixture
def experiments(client):
    exp1 = client.create_experiment("test_exp_1")
    exp2 = client.create_experiment("test_exp_2")
    exp3 = client.create_experiment("test_exp_3")
    return [exp1, exp2, exp3]


@pytest.fixture
def experiment(client):
    return client.create_experiment("test_trace_experiment")


def test_create_dataset(mock_client):
    expected_dataset = EntityEvaluationDataset(
        dataset_id="test_id",
        name="test_dataset",
        digest="abc123",
        created_time=123456789,
        last_update_time=123456789,
        tags={"environment": "production", "version": "1.0"},
    )
    mock_client.create_dataset.return_value = expected_dataset

    result = create_dataset(
        name="test_dataset",
        experiment_id=["exp1", "exp2"],
        tags={"environment": "production", "version": "1.0"},
    )

    assert result == expected_dataset
    mock_client.create_dataset.assert_called_once_with(
        name="test_dataset",
        experiment_id=["exp1", "exp2"],
        tags={"environment": "production", "version": "1.0"},
    )


def test_create_dataset_single_experiment_id(mock_client):
    expected_dataset = EntityEvaluationDataset(
        dataset_id="test_id",
        name="test_dataset",
        digest="abc123",
        created_time=123456789,
        last_update_time=123456789,
    )
    mock_client.create_dataset.return_value = expected_dataset

    result = create_dataset(
        name="test_dataset",
        experiment_id="exp1",
    )

    assert result == expected_dataset
    mock_client.create_dataset.assert_called_once_with(
        name="test_dataset",
        experiment_id=["exp1"],
        tags=None,
    )


def test_create_dataset_with_empty_tags(mock_client):
    expected_dataset = EntityEvaluationDataset(
        dataset_id="test_id",
        name="test_dataset",
        digest="abc123",
        created_time=123456789,
        last_update_time=123456789,
        tags={},
    )
    mock_client.create_dataset.return_value = expected_dataset

    result = create_dataset(
        name="test_dataset",
        experiment_id=["exp1"],
        tags={},
    )

    assert result == expected_dataset
    mock_client.create_dataset.assert_called_once_with(
        name="test_dataset",
        experiment_id=["exp1"],
        tags={},
    )


def test_create_dataset_databricks(mock_databricks_environment):
    mock_dataset = mock.Mock()
    with mock.patch.dict(
        "sys.modules",
        {
            "databricks.agents.datasets": mock.Mock(
                create_dataset=mock.Mock(return_value=mock_dataset)
            )
        },
    ):
        result = create_dataset(
            name="catalog.schema.table",
            experiment_id=["exp1", "exp2"],
        )

        sys.modules["databricks.agents.datasets"].create_dataset.assert_called_once_with(
            "catalog.schema.table", ["exp1", "exp2"]
        )
        assert isinstance(result, EvaluationDataset)


def test_get_dataset(mock_client):
    expected_dataset = EntityEvaluationDataset(
        dataset_id="test_id",
        name="test_dataset",
        digest="abc123",
        created_time=123456789,
        last_update_time=123456789,
    )
    mock_client.get_dataset.return_value = expected_dataset

    result = get_dataset(dataset_id="test_id")

    assert result == expected_dataset
    mock_client.get_dataset.assert_called_once_with("test_id")


def test_get_dataset_missing_id():
    with pytest.raises(ValueError, match="Either 'name' or 'dataset_id' must be provided"):
        get_dataset()


def test_get_dataset_databricks(mock_databricks_environment):
    mock_dataset = mock.Mock()
    with mock.patch.dict(
        "sys.modules",
        {"databricks.agents.datasets": mock.Mock(get_dataset=mock.Mock(return_value=mock_dataset))},
    ):
        result = get_dataset(name="catalog.schema.table")

        sys.modules["databricks.agents.datasets"].get_dataset.assert_called_once_with(
            "catalog.schema.table"
        )
        assert isinstance(result, EvaluationDataset)


def test_get_dataset_databricks_missing_name(mock_databricks_environment):
    with pytest.raises(ValueError, match="Parameter 'name' is required"):
        get_dataset(dataset_id="test_id")


def test_get_dataset_by_name_oss(experiments):
    dataset = create_dataset(
        name="unique_dataset_name",
        experiment_id=experiments[0],
        tags={"test": "get_by_name"},
    )

    retrieved = get_dataset(name="unique_dataset_name")

    assert retrieved.dataset_id == dataset.dataset_id
    assert retrieved.name == "unique_dataset_name"
    assert retrieved.tags["test"] == "get_by_name"


def test_get_dataset_by_name_not_found(client):
    with pytest.raises(MlflowException, match="Dataset with name 'nonexistent_dataset' not found"):
        get_dataset(name="nonexistent_dataset")


def test_get_dataset_by_name_multiple_matches(experiments):
    create_dataset(
        name="duplicate_name",
        experiment_id=experiments[0],
    )
    create_dataset(
        name="duplicate_name",
        experiment_id=experiments[1],
    )

    with pytest.raises(MlflowException, match="Multiple datasets found with name 'duplicate_name'"):
        get_dataset(name="duplicate_name")


def test_get_dataset_both_name_and_id_error(experiments):
    dataset = create_dataset(
        name="test_dataset_both",
        experiment_id=experiments[0],
    )

    with pytest.raises(ValueError, match="Cannot specify both 'name' and 'dataset_id'"):
        get_dataset(name="test_dataset_both", dataset_id=dataset.dataset_id)


def test_get_dataset_neither_name_nor_id_error(client):
    with pytest.raises(ValueError, match="Either 'name' or 'dataset_id' must be provided"):
        get_dataset()


@pytest.mark.parametrize(
    "name",
    [
        "dataset's_with_single_quote",
        'dataset"with_double_quote',
    ],
)
def test_get_dataset_name_with_quotes(experiments, name):
    dataset = create_dataset(name=name, experiment_id=experiments[0])

    retrieved = get_dataset(name=name)

    assert retrieved.dataset_id == dataset.dataset_id
    assert retrieved.name == name


def test_delete_dataset(mock_client):
    delete_dataset(dataset_id="test_id")

    mock_client.delete_dataset.assert_called_once_with("test_id")


def test_delete_dataset_missing_id():
    with pytest.raises(ValueError, match="Parameter 'dataset_id' is required"):
        delete_dataset()


def test_delete_dataset_databricks(mock_databricks_environment):
    with mock.patch.dict(
        "sys.modules",
        {"databricks.agents.datasets": mock.Mock(delete_dataset=mock.Mock())},
    ):
        delete_dataset(name="catalog.schema.table")

        sys.modules["databricks.agents.datasets"].delete_dataset.assert_called_once_with(
            "catalog.schema.table"
        )


def test_search_datasets_with_mock(mock_client):
    datasets = [
        EntityEvaluationDataset(
            dataset_id="id1",
            name="dataset1",
            digest="digest1",
            created_time=123456789,
            last_update_time=123456789,
        ),
        EntityEvaluationDataset(
            dataset_id="id2",
            name="dataset2",
            digest="digest2",
            created_time=123456789,
            last_update_time=123456789,
        ),
    ]
    mock_client.search_datasets.return_value = PagedList(datasets, None)

    result = search_datasets(
        experiment_ids=["exp1", "exp2"],
        filter_string="name LIKE 'test%'",
        max_results=100,
        order_by=["created_time DESC"],
    )

    assert len(result) == 2
    assert isinstance(result, list)

    mock_client.search_datasets.assert_called_once_with(
        experiment_ids=["exp1", "exp2"],
        filter_string="name LIKE 'test%'",
        max_results=50,
        order_by=["created_time DESC"],
        page_token=None,
    )


def test_search_datasets_single_experiment_id(mock_client):
    datasets = [
        EntityEvaluationDataset(
            dataset_id="id1",
            name="dataset1",
            digest="digest1",
            created_time=123456789,
            last_update_time=123456789,
        )
    ]
    mock_client.search_datasets.return_value = PagedList(datasets, None)

    # When no max_results is specified, it defaults to None which means get all
    # Mock time to have a consistent filter_string
    with mock.patch("time.time", return_value=1234567890):
        search_datasets(experiment_ids="exp1")

    # The pagination wrapper will use SEARCH_EVALUATION_DATASETS_MAX_RESULTS as the page size
    # Now the function adds default filter (last 7 days) and order_by when not specified
    seven_days_ago = int((1234567890 - 7 * 24 * 60 * 60) * 1000)
    mock_client.search_datasets.assert_called_once_with(
        experiment_ids=["exp1"],
        filter_string=f"created_time >= {seven_days_ago}",
        max_results=SEARCH_EVALUATION_DATASETS_MAX_RESULTS,  # Page size
        order_by=["created_time DESC"],
        page_token=None,
    )


def test_search_datasets_pagination_handling(mock_client):
    page1_datasets = [
        EntityEvaluationDataset(
            dataset_id=f"id{i}",
            name=f"dataset{i}",
            digest=f"digest{i}",
            created_time=123456789,
            last_update_time=123456789,
        )
        for i in range(3)
    ]

    page2_datasets = [
        EntityEvaluationDataset(
            dataset_id=f"id{i}",
            name=f"dataset{i}",
            digest=f"digest{i}",
            created_time=123456789,
            last_update_time=123456789,
        )
        for i in range(3, 5)
    ]

    mock_client.search_datasets.side_effect = [
        PagedList(page1_datasets, "token1"),
        PagedList(page2_datasets, None),
    ]

    result = search_datasets(experiment_ids=["exp1"], max_results=10)

    assert len(result) == 5
    assert isinstance(result, list)

    assert mock_client.search_datasets.call_count == 2

    first_call = mock_client.search_datasets.call_args_list[0]
    assert first_call[1]["page_token"] is None

    second_call = mock_client.search_datasets.call_args_list[1]
    assert second_call[1]["page_token"] == "token1"


def test_search_datasets_single_page(mock_client):
    datasets = [
        EntityEvaluationDataset(
            dataset_id="id1",
            name="dataset1",
            digest="digest1",
            created_time=123456789,
            last_update_time=123456789,
        )
    ]

    mock_client.search_datasets.return_value = PagedList(datasets, None)

    result = search_datasets(max_results=10)

    assert len(result) == 1
    assert isinstance(result, list)

    assert mock_client.search_datasets.call_count == 1


def test_search_datasets_databricks(mock_databricks_environment, mock_client):
    datasets = [
        EntityEvaluationDataset(
            dataset_id="id1",
            name="dataset1",
            digest="digest1",
            created_time=123456789,
            last_update_time=123456789,
        ),
    ]
    mock_client.search_datasets.return_value = PagedList(datasets, None)

    result = search_datasets(experiment_ids=["exp1"])

    assert len(result) == 1
    assert isinstance(result, list)

    # Verify that default filter_string and order_by are NOT set for Databricks
    # (since these parameters may not be supported by all Databricks backends)
    mock_client.search_datasets.assert_called_once()
    call_kwargs = mock_client.search_datasets.call_args.kwargs
    assert call_kwargs.get("filter_string") is None
    assert call_kwargs.get("order_by") is None


def test_databricks_import_error():
    with (
        mock.patch("mlflow.genai.datasets.is_databricks_uri", return_value=True),
        mock.patch.dict("sys.modules", {"databricks.agents.datasets": None}),
        mock.patch("builtins.__import__", side_effect=ImportError("No module")),
    ):
        with pytest.raises(ImportError, match="databricks-agents"):
            create_dataset(name="test", experiment_id="exp1")


def test_databricks_profile_uri_support():
    mock_dataset = mock.Mock()
    with (
        mock.patch(
            "mlflow.genai.datasets.get_tracking_uri",
            return_value="databricks://profilename",
        ),
        mock.patch.dict(
            "sys.modules",
            {
                "databricks.agents.datasets": mock.Mock(
                    get_dataset=mock.Mock(return_value=mock_dataset),
                    create_dataset=mock.Mock(return_value=mock_dataset),
                    delete_dataset=mock.Mock(),
                )
            },
        ),
    ):
        result = get_dataset(name="catalog.schema.table")
        sys.modules["databricks.agents.datasets"].get_dataset.assert_called_once_with(
            "catalog.schema.table"
        )
        assert isinstance(result, EvaluationDataset)

        result2 = create_dataset(name="catalog.schema.table2", experiment_id=["exp1"])
        sys.modules["databricks.agents.datasets"].create_dataset.assert_called_once_with(
            "catalog.schema.table2", ["exp1"]
        )
        assert isinstance(result2, EvaluationDataset)

        delete_dataset(name="catalog.schema.table3")
        sys.modules["databricks.agents.datasets"].delete_dataset.assert_called_once_with(
            "catalog.schema.table3"
        )


def test_databricks_profile_env_var_set_from_uri(monkeypatch):
    mock_dataset = mock.Mock()
    profile_values_during_calls = []

    def mock_get_dataset(name):
        profile_values_during_calls.append((
            "get_dataset",
            os.environ.get("DATABRICKS_CONFIG_PROFILE"),
        ))
        return mock_dataset

    def mock_create_dataset(name, experiment_ids):
        profile_values_during_calls.append((
            "create_dataset",
            os.environ.get("DATABRICKS_CONFIG_PROFILE"),
        ))
        return mock_dataset

    def mock_delete_dataset(name):
        profile_values_during_calls.append((
            "delete_dataset",
            os.environ.get("DATABRICKS_CONFIG_PROFILE"),
        ))

    mock_agents_module = mock.Mock(
        get_dataset=mock_get_dataset,
        create_dataset=mock_create_dataset,
        delete_dataset=mock_delete_dataset,
    )
    monkeypatch.setitem(sys.modules, "databricks.agents.datasets", mock_agents_module)
    monkeypatch.setattr("mlflow.genai.datasets.get_tracking_uri", lambda: "databricks://myprofile")

    assert "DATABRICKS_CONFIG_PROFILE" not in os.environ

    get_dataset(name="catalog.schema.table")
    create_dataset(name="catalog.schema.table", experiment_id="exp1")
    delete_dataset(name="catalog.schema.table")

    assert "DATABRICKS_CONFIG_PROFILE" not in os.environ

    assert profile_values_during_calls == [
        ("get_dataset", "myprofile"),
        ("create_dataset", "myprofile"),
        ("delete_dataset", "myprofile"),
    ]


def test_databricks_profile_env_var_overridden_and_restored(monkeypatch):
    mock_dataset = mock.Mock()
    profile_during_call = None

    def mock_get_dataset(name):
        nonlocal profile_during_call
        profile_during_call = os.environ.get("DATABRICKS_CONFIG_PROFILE")
        return mock_dataset

    mock_agents_module = mock.Mock(get_dataset=mock_get_dataset)
    monkeypatch.setitem(sys.modules, "databricks.agents.datasets", mock_agents_module)
    monkeypatch.setattr("mlflow.genai.datasets.get_tracking_uri", lambda: "databricks://myprofile")
    monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", "original_profile")

    assert os.environ.get("DATABRICKS_CONFIG_PROFILE") == "original_profile"

    get_dataset(name="catalog.schema.table")

    assert os.environ.get("DATABRICKS_CONFIG_PROFILE") == "original_profile"
    assert profile_during_call == "myprofile"


def test_databricks_dataset_merge_records_uses_profile(monkeypatch):
    profile_during_merge = None
    profile_during_to_df = None

    mock_inner_dataset = mock.Mock()
    mock_inner_dataset.digest = "test_digest"
    mock_inner_dataset.name = "catalog.schema.table"
    mock_inner_dataset.dataset_id = "dataset-123"

    def mock_merge_records(records):
        nonlocal profile_during_merge
        profile_during_merge = os.environ.get("DATABRICKS_CONFIG_PROFILE")
        return mock_inner_dataset

    def mock_to_df():
        nonlocal profile_during_to_df
        profile_during_to_df = os.environ.get("DATABRICKS_CONFIG_PROFILE")
        import pandas as pd

        return pd.DataFrame({"test": [1, 2, 3]})

    mock_inner_dataset.merge_records = mock_merge_records
    mock_inner_dataset.to_df = mock_to_df

    def mock_get_dataset(name):
        return mock_inner_dataset

    mock_agents_module = mock.Mock(get_dataset=mock_get_dataset)
    monkeypatch.setitem(sys.modules, "databricks.agents.datasets", mock_agents_module)
    monkeypatch.setattr("mlflow.genai.datasets.get_tracking_uri", lambda: "databricks://myprofile")

    assert "DATABRICKS_CONFIG_PROFILE" not in os.environ

    dataset = get_dataset(name="catalog.schema.table")

    assert "DATABRICKS_CONFIG_PROFILE" not in os.environ

    dataset.merge_records([{"inputs": {"q": "test"}}])
    assert profile_during_merge == "myprofile"
    assert "DATABRICKS_CONFIG_PROFILE" not in os.environ

    dataset.to_df()
    assert profile_during_to_df == "myprofile"
    assert "DATABRICKS_CONFIG_PROFILE" not in os.environ


def test_create_dataset_with_user_tag(experiments):
    dataset = create_dataset(
        name="test_user_attribution",
        experiment_id=experiments[0],
        tags={"environment": "test", MLFLOW_USER: "john_doe"},
    )

    assert dataset.name == "test_user_attribution"
    assert dataset.tags[MLFLOW_USER] == "john_doe"
    assert dataset.created_by == "john_doe"

    dataset2 = create_dataset(
        name="test_no_user",
        experiment_id=experiments[0],
        tags={"environment": "test"},
    )

    assert dataset2.name == "test_no_user"
    assert isinstance(dataset2.tags[MLFLOW_USER], str)
    assert dataset2.created_by == dataset2.tags[MLFLOW_USER]


def test_create_and_get_dataset(experiments):
    dataset = create_dataset(
        name="qa_evaluation_v1",
        experiment_id=[experiments[0], experiments[1]],
        tags={"source": "manual_curation", "environment": "test"},
    )

    assert dataset.name == "qa_evaluation_v1"
    assert dataset.tags["source"] == "manual_curation"
    assert dataset.tags["environment"] == "test"
    assert len(dataset.experiment_ids) == 2
    assert dataset.dataset_id is not None

    retrieved = get_dataset(dataset_id=dataset.dataset_id)

    assert retrieved.dataset_id == dataset.dataset_id
    assert retrieved.name == dataset.name
    assert retrieved.tags == dataset.tags
    assert set(retrieved.experiment_ids) == {experiments[0], experiments[1]}


def test_create_dataset_minimal_params(client):
    dataset = create_dataset(name="minimal_dataset")

    assert dataset.name == "minimal_dataset"
    assert "mlflow.user" not in dataset.tags or isinstance(dataset.tags.get("mlflow.user"), str)
    assert dataset.experiment_ids == ["0"]


def test_active_record_pattern_merge_records(experiments):
    dataset = create_dataset(
        name="active_record_test",
        experiment_id=experiments[0],
    )

    records_batch1 = [
        {
            "inputs": {"question": "What is MLflow?"},
            "outputs": {
                "answer": "MLflow is an open source platform for managing the ML lifecycle",
                "key1": "value1",
            },
            "expectations": {
                "answer": "MLflow is an open source platform",
                "key2": "value2",
            },
            "tags": {"difficulty": "easy"},
        },
        {
            "inputs": {"question": "What is Python?"},
            "outputs": {"answer": "Python is a versatile programming language"},
            "expectations": {"answer": "Python is a programming language"},
            "tags": {"difficulty": "easy"},
        },
    ]

    records_batch2 = [
        {
            "inputs": {"question": "What is MLflow?"},
            "outputs": {"answer": "MLflow is a popular ML lifecycle platform"},
            "expectations": {"answer": "MLflow is an ML lifecycle platform"},
            "tags": {"category": "ml"},
        },
        {
            "inputs": {"question": "What is Docker?"},
            "outputs": {"answer": "Docker is a popular containerization platform"},
            "expectations": {"answer": "Docker is a containerization platform"},
            "tags": {"difficulty": "medium"},
        },
    ]

    dataset.merge_records(records_batch1)

    df1 = dataset.to_df()
    assert len(df1) == 2

    mlflow_record = df1[df1["inputs"].apply(lambda x: x.get("question") == "What is MLflow?")].iloc[
        0
    ]
    assert mlflow_record["expectations"] == {
        "answer": "MLflow is an open source platform",
        "key2": "value2",
    }
    assert mlflow_record["outputs"] == {
        "answer": "MLflow is an open source platform for managing the ML lifecycle",
        "key1": "value1",
    }
    assert mlflow_record["tags"]["difficulty"] == "easy"
    assert "category" not in mlflow_record["tags"]

    dataset.merge_records(records_batch2)

    df2 = dataset.to_df()
    assert len(df2) == 3

    mlflow_record_updated = df2[
        df2["inputs"].apply(lambda x: x.get("question") == "What is MLflow?")
    ].iloc[0]

    assert mlflow_record_updated["expectations"] == {
        "answer": "MLflow is an ML lifecycle platform",
        "key2": "value2",
    }
    assert mlflow_record_updated["outputs"] == {
        "answer": "MLflow is a popular ML lifecycle platform"
    }
    assert mlflow_record_updated["tags"]["difficulty"] == "easy"
    assert mlflow_record_updated["tags"]["category"] == "ml"

    # Verify that the new Docker record also has outputs
    docker_record = df2[df2["inputs"].apply(lambda x: x.get("question") == "What is Docker?")].iloc[
        0
    ]
    assert docker_record["outputs"]["answer"] == "Docker is a popular containerization platform"
    assert docker_record["expectations"]["answer"] == "Docker is a containerization platform"
    assert docker_record["tags"]["difficulty"] == "medium"


def test_dataset_with_dataframe_records(experiments):
    dataset = create_dataset(
        name="dataframe_test",
        experiment_id=experiments[0],
        tags={"source": "csv", "file": "test_data.csv"},
    )

    df = pd.DataFrame([
        {
            "inputs": {"text": "The movie was amazing!", "model": "sentiment-v1"},
            "expectations": {"sentiment": "positive", "confidence": 0.95},
            "tags": {"source": "imdb"},
        },
        {
            "inputs": {"text": "Terrible experience", "model": "sentiment-v1"},
            "expectations": {"sentiment": "negative", "confidence": 0.88},
            "tags": {"source": "yelp"},
        },
    ])

    dataset.merge_records(df)

    result_df = dataset.to_df()
    assert len(result_df) == 2
    assert all(col in result_df.columns for col in ["inputs", "expectations", "tags"])

    # Check that all expected records are present (order-agnostic)
    texts = {record["inputs"]["text"] for _, record in result_df.iterrows()}
    expected_texts = {"The movie was amazing!", "Terrible experience"}
    assert texts == expected_texts

    sentiments = {record["expectations"]["sentiment"] for _, record in result_df.iterrows()}
    expected_sentiments = {"positive", "negative"}
    assert sentiments == expected_sentiments


def test_search_datasets(experiments):
    for i in range(5):
        create_dataset(
            name=f"search_test_{i}",
            experiment_id=[experiments[i % len(experiments)]],
            tags={"type": "human" if i % 2 == 0 else "trace", "index": str(i)},
        )

    all_results = search_datasets()
    assert len(all_results) == 5

    exp0_results = search_datasets(experiment_ids=experiments[0])
    assert len(exp0_results) == 2

    human_results = search_datasets(filter_string="name LIKE 'search_test_%'")
    assert len(human_results) == 5

    limited_results = search_datasets(max_results=2)
    assert len(limited_results) == 2

    more_results = search_datasets(max_results=4)
    assert len(more_results) == 4


def test_delete_dataset(experiments):
    dataset = create_dataset(
        name="to_be_deleted",
        experiment_id=[experiments[0], experiments[1]],
        tags={"env": "test", "version": "1.0"},
    )
    dataset_id = dataset.dataset_id

    dataset.merge_records([{"inputs": {"q": "test"}, "expectations": {"a": "answer"}}])

    retrieved = get_dataset(dataset_id=dataset_id)
    assert retrieved is not None
    assert len(retrieved.to_df()) == 1

    delete_dataset(dataset_id=dataset_id)

    with pytest.raises(MlflowException, match="Could not find|not found"):
        get_dataset(dataset_id=dataset_id)

    search_results = search_datasets(experiment_ids=[experiments[0], experiments[1]])
    found_ids = [d.dataset_id for d in search_results]
    assert dataset_id not in found_ids


def test_dataset_lifecycle_workflow(experiments):
    dataset = create_dataset(
        name="qa_eval_prod_v1",
        experiment_id=[experiments[0], experiments[1]],
        tags={"source": "qa_team_annotations", "team": "qa", "env": "prod"},
    )

    initial_cases = [
        {
            "inputs": {"question": "What is the capital of France?"},
            "expectations": {"answer": "Paris", "confidence": "high"},
            "tags": {"category": "geography", "difficulty": "easy"},
        },
        {
            "inputs": {"question": "Explain quantum computing"},
            "expectations": {"answer": "Quantum computing uses quantum mechanics principles"},
            "tags": {"category": "science", "difficulty": "hard"},
        },
    ]
    dataset.merge_records(initial_cases)

    dataset_id = dataset.dataset_id
    retrieved = get_dataset(dataset_id=dataset_id)
    df = retrieved.to_df()
    assert len(df) == 2

    additional_cases = [
        {
            "inputs": {"question": "What is 2+2?"},
            "expectations": {"answer": "4", "confidence": "high"},
            "tags": {"category": "math", "difficulty": "easy"},
        },
    ]
    retrieved.merge_records(additional_cases)

    found = search_datasets(
        experiment_ids=experiments[0],
        filter_string="name LIKE 'qa_eval%'",
    )
    assert len(found) == 1
    assert found[0].dataset_id == dataset_id

    final_dataset = get_dataset(dataset_id=dataset_id)
    final_df = final_dataset.to_df()
    assert len(final_df) == 3

    categories = set()
    for _, row in final_df.iterrows():
        if row["tags"] and "category" in row["tags"]:
            categories.add(row["tags"]["category"])
    assert categories == {"geography", "science", "math"}


def test_error_handling_filestore_backend(tmp_path):
    pytest.skip("FileStore is no longer supported.")
    file_uri = f"file://{tmp_path}"
    mlflow.set_tracking_uri(file_uri)

    with pytest.raises(MlflowException, match="not supported with FileStore") as exc:
        create_dataset(name="test")
    assert exc.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="not supported with FileStore") as exc:
        get_dataset(dataset_id="test_id")
    assert exc.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="not supported with FileStore") as exc:
        search_datasets()
    assert exc.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="not supported with FileStore") as exc:
        delete_dataset(dataset_id="test_id")
    assert exc.value.error_code == "FEATURE_DISABLED"


def test_single_experiment_id_handling(experiments):
    dataset = create_dataset(
        name="single_exp_test",
        experiment_id=experiments[0],
    )

    assert isinstance(dataset.experiment_ids, list)
    assert dataset.experiment_ids == [experiments[0]]

    results = search_datasets(experiment_ids=experiments[0])
    found_ids = [d.dataset_id for d in results]
    assert dataset.dataset_id in found_ids


def test_trace_to_evaluation_dataset_integration(experiments):
    trace_inputs = [
        {"question": "What is MLflow?", "context": "ML platforms"},
        {"question": "What is Python?", "context": "programming"},
        {"question": "What is MLflow?", "context": "ML platforms"},
    ]

    created_trace_ids = []
    for i, inputs in enumerate(trace_inputs):
        with mlflow.start_run(experiment_id=experiments[i % 2]):
            with mlflow.start_span(name=f"qa_trace_{i}") as span:
                span.set_inputs(inputs)
                span.set_outputs({"answer": f"Answer for {inputs['question']}"})
                span.set_attributes({"model": "test-model", "temperature": "0.7"})
                trace_id = span.trace_id
                created_trace_ids.append(trace_id)

                mlflow.log_expectation(
                    trace_id=trace_id,
                    name="expected_answer",
                    value=f"Detailed answer for {inputs['question']}",
                )
                mlflow.log_expectation(
                    trace_id=trace_id,
                    name="quality_score",
                    value=0.85 + i * 0.05,
                )

    traces = mlflow.search_traces(
        locations=[experiments[0], experiments[1]],
        max_results=10,
        return_type="list",
    )
    assert len(traces) == 3

    dataset = create_dataset(
        name="trace_eval_dataset",
        experiment_id=[experiments[0], experiments[1]],
        tags={"source": "test_traces", "type": "trace_integration"},
    )

    dataset.merge_records(traces)

    df = dataset.to_df()
    assert len(df) == 2

    for _, record in df.iterrows():
        assert "inputs" in record
        assert "question" in record["inputs"]
        assert "context" in record["inputs"]
        assert record.get("source_type") == "TRACE"
        assert record.get("source_id") is not None

    mlflow_records = df[df["inputs"].apply(lambda x: x.get("question") == "What is MLflow?")]
    assert len(mlflow_records) == 1

    with mlflow.start_run(experiment_id=experiments[0]):
        with mlflow.start_span(name="additional_trace") as span:
            span.set_inputs({"question": "What is Docker?", "context": "containers"})
            span.set_outputs({"answer": "Docker is a containerization platform"})
            span.set_attributes({"model": "test-model"})

    all_traces = mlflow.search_traces(
        locations=[experiments[0], experiments[1]], max_results=10, return_type="list"
    )
    assert len(all_traces) == 4

    new_trace = None
    for trace in all_traces:
        root_span = trace.data._get_root_span() if hasattr(trace, "data") else None
        if root_span and root_span.inputs and root_span.inputs.get("question") == "What is Docker?":
            new_trace = trace
            break

    assert new_trace is not None

    dataset.merge_records([new_trace])

    final_df = dataset.to_df()
    assert len(final_df) == 3

    retrieved = get_dataset(dataset_id=dataset.dataset_id)
    retrieved_df = retrieved.to_df()
    assert len(retrieved_df) == 3

    delete_dataset(dataset_id=dataset.dataset_id)

    with pytest.raises(MlflowException, match="Could not find|not found"):
        get_dataset(dataset_id=dataset.dataset_id)

    search_results = search_datasets(
        experiment_ids=[experiments[0], experiments[1]], max_results=100
    )
    found_dataset_ids = [d.dataset_id for d in search_results]
    assert dataset.dataset_id not in found_dataset_ids

    all_datasets = search_datasets(max_results=100)
    all_dataset_ids = [d.dataset_id for d in all_datasets]
    assert dataset.dataset_id not in all_dataset_ids


def test_search_traces_dataframe_to_dataset_integration(experiments):
    for i in range(3):
        with mlflow.start_run(experiment_id=experiments[0]):
            with mlflow.start_span(name=f"test_span_{i}") as span:
                span.set_inputs({"question": f"Question {i}?", "temperature": 0.7})
                span.set_outputs({"answer": f"Answer {i}"})

                mlflow.log_expectation(
                    trace_id=span.trace_id,
                    name="expected_answer",
                    value=f"Expected answer {i}",
                )
                mlflow.log_expectation(
                    trace_id=span.trace_id,
                    name="min_score",
                    value=0.8,
                )

    traces_df = mlflow.search_traces(
        locations=[experiments[0]],
    )

    assert "trace" in traces_df.columns
    assert "assessments" in traces_df.columns
    assert len(traces_df) == 3

    dataset = create_dataset(
        name="traces_dataframe_dataset",
        experiment_id=experiments[0],
        tags={"source": "search_traces", "format": "dataframe"},
    )

    dataset.merge_records(traces_df)

    result_df = dataset.to_df()
    assert len(result_df) == 3

    for idx, row in result_df.iterrows():
        assert "inputs" in row
        assert "expectations" in row
        assert "source_type" in row
        assert row["source_type"] == "TRACE"

        assert "question" in row["inputs"]
        question_text = row["inputs"]["question"]
        assert question_text.startswith("Question ")
        assert question_text.endswith("?")
        question_num = int(question_text.replace("Question ", "").replace("?", ""))
        assert 0 <= question_num <= 2

        assert "expected_answer" in row["expectations"]
        assert f"Expected answer {question_num}" == row["expectations"]["expected_answer"]
        assert "min_score" in row["expectations"]
        assert row["expectations"]["min_score"] == 0.8


def test_trace_to_dataset_with_assessments(client, experiment):
    trace_data = [
        {
            "inputs": {"question": "What is MLflow?", "context": "ML platforms"},
            "outputs": {"answer": "MLflow is an open source platform for ML lifecycle"},
            "expectations": {
                "correctness": True,
                "completeness": 0.8,
            },
        },
        {
            "inputs": {
                "question": "What is Python?",
                "context": "programming languages",
            },
            "outputs": {"answer": "Python is a high-level programming language"},
            "expectations": {
                "correctness": True,
            },
        },
        {
            "inputs": {"question": "What is Docker?", "context": "containerization"},
            "outputs": {"answer": "Docker is a container platform"},
            "expectations": {},
        },
    ]

    created_traces = []
    for i, data in enumerate(trace_data):
        with mlflow.start_run(experiment_id=experiment):
            with mlflow.start_span(name=f"qa_trace_{i}") as span:
                span.set_inputs(data["inputs"])
                span.set_outputs(data["outputs"])
                span.set_attributes({"model": "test-model", "temperature": 0.7})
                trace_id = span.trace_id

                for name, value in data["expectations"].items():
                    mlflow.log_expectation(
                        trace_id=trace_id,
                        name=name,
                        value=value,
                        span_id=span.span_id,
                    )

        trace = client.get_trace(trace_id)
        created_traces.append(trace)

    dataset = create_dataset(
        name="trace_assessment_dataset",
        experiment_id=[experiment],
        tags={"source": "trace_integration_test", "version": "1.0"},
    )

    dataset.merge_records(created_traces)

    df = dataset.to_df()
    assert len(df) == 3

    mlflow_record = df[df["inputs"].apply(lambda x: x.get("question") == "What is MLflow?")].iloc[0]
    assert mlflow_record["inputs"]["question"] == "What is MLflow?"
    assert mlflow_record["inputs"]["context"] == "ML platforms"

    assert "expectations" in mlflow_record
    assert mlflow_record["expectations"]["correctness"] is True
    assert mlflow_record["expectations"]["completeness"] == 0.8

    assert mlflow_record["source_type"] == "TRACE"
    assert mlflow_record["source_id"] is not None

    python_record = df[df["inputs"].apply(lambda x: x.get("question") == "What is Python?")].iloc[0]
    assert python_record["expectations"]["correctness"] is True
    assert len(python_record["expectations"]) == 1

    docker_record = df[df["inputs"].apply(lambda x: x.get("question") == "What is Docker?")].iloc[0]
    assert docker_record["expectations"] is None or docker_record["expectations"] == {}

    retrieved = get_dataset(dataset_id=dataset.dataset_id)
    assert retrieved.tags["source"] == "trace_integration_test"
    assert retrieved.tags["version"] == "1.0"
    assert set(retrieved.experiment_ids) == {experiment}


def test_trace_deduplication_with_assessments(client, experiment):
    trace_ids = []
    for i in range(3):
        with mlflow.start_run(experiment_id=experiment):
            with mlflow.start_span(name=f"duplicate_trace_{i}") as span:
                span.set_inputs({"question": "What is AI?", "model": "gpt-4"})
                span.set_outputs({"answer": f"AI is artificial intelligence (version {i})"})
                trace_id = span.trace_id
                trace_ids.append(trace_id)

                mlflow.log_expectation(
                    trace_id=trace_id,
                    name="quality",
                    value=0.5 + i * 0.2,
                    span_id=span.span_id,
                )

    traces = [client.get_trace(tid) for tid in trace_ids]

    dataset = create_dataset(
        name="dedup_test",
        experiment_id=experiment,
        tags={"test": "deduplication"},
    )
    dataset.merge_records(traces)

    df = dataset.to_df()
    assert len(df) == 1

    record = df.iloc[0]
    assert record["inputs"]["question"] == "What is AI?"
    assert record["expectations"]["quality"] == 0.9
    assert record["source_id"] in trace_ids


def test_mixed_record_types_with_traces(client, experiment):
    with mlflow.start_run(experiment_id=experiment):
        with mlflow.start_span(name="mixed_test_trace") as span:
            span.set_inputs({"question": "What is ML?", "context": "machine learning"})
            span.set_outputs({"answer": "ML stands for Machine Learning"})
            trace_id = span.trace_id

            mlflow.log_expectation(
                trace_id=trace_id,
                name="accuracy",
                value=0.95,
                span_id=span.span_id,
            )

    trace = client.get_trace(trace_id)

    dataset = create_dataset(
        name="mixed_records_test",
        experiment_id=experiment,
        tags={"type": "mixed", "test": "true"},
    )

    manual_records = [
        {
            "inputs": {"question": "What is AI?"},
            "expectations": {"correctness": True},
            "tags": {"source": "manual"},
        },
        {
            "inputs": {"question": "What is Python?"},
            "expectations": {"correctness": True},
            "tags": {"source": "manual"},
        },
    ]
    dataset.merge_records(manual_records)

    df1 = dataset.to_df()
    assert len(df1) == 2

    dataset.merge_records([trace])

    df2 = dataset.to_df()
    assert len(df2) == 3

    ml_record = df2[df2["inputs"].apply(lambda x: x.get("question") == "What is ML?")].iloc[0]
    assert ml_record["expectations"]["accuracy"] == 0.95
    assert ml_record["source_type"] == "TRACE"

    manual_questions = {"What is AI?", "What is Python?"}
    manual_records_df = df2[df2["inputs"].apply(lambda x: x.get("question") in manual_questions)]
    assert len(manual_records_df) == 2

    for _, record in manual_records_df.iterrows():
        assert record.get("source_type") != "TRACE"


def test_trace_without_root_span_inputs(client, experiment):
    with mlflow.start_run(experiment_id=experiment):
        with mlflow.start_span(name="no_inputs_trace") as span:
            span.set_outputs({"result": "some output"})
            trace_id = span.trace_id

    trace = client.get_trace(trace_id)

    dataset = create_dataset(
        name="no_inputs_test",
        experiment_id=experiment,
    )

    dataset.merge_records([trace])

    df = dataset.to_df()
    assert len(df) == 1
    assert df.iloc[0]["inputs"] == {}
    assert df.iloc[0]["expectations"] is None or df.iloc[0]["expectations"] == {}


def test_error_handling_invalid_trace_types(client, experiment):
    dataset = create_dataset(
        name="error_test",
        experiment_id=experiment,
    )

    with mlflow.start_run(experiment_id=experiment):
        with mlflow.start_span(name="valid_trace") as span:
            span.set_inputs({"q": "test"})
            trace_id = span.trace_id

    valid_trace = client.get_trace(trace_id)

    with pytest.raises(MlflowException, match="Mixed types in trace list"):
        dataset.merge_records([valid_trace, {"inputs": {"q": "dict record"}}])

    with pytest.raises(MlflowException, match="Mixed types in trace list"):
        dataset.merge_records([valid_trace, "not a trace"])


def test_trace_integration_end_to_end(client, experiment):
    traces_to_create = [
        {
            "name": "successful_qa",
            "inputs": {"question": "What is the capital of France?", "language": "en"},
            "outputs": {"answer": "Paris", "confidence": 0.99},
            "expectations": {"correctness": True, "confidence_threshold": 0.8},
        },
        {
            "name": "incorrect_qa",
            "inputs": {"question": "What is 2+2?", "language": "en"},
            "outputs": {"answer": "5", "confidence": 0.5},
            "expectations": {"correctness": False},
        },
        {
            "name": "multilingual_qa",
            "inputs": {"question": "¿Cómo estás?", "language": "es"},
            "outputs": {"answer": "I'm doing well, thank you!", "confidence": 0.9},
            "expectations": {"language_match": False, "politeness": True},
        },
    ]

    created_trace_ids = []
    for trace_config in traces_to_create:
        with mlflow.start_run(experiment_id=experiment):
            with mlflow.start_span(name=trace_config["name"]) as span:
                span.set_inputs(trace_config["inputs"])
                span.set_outputs(trace_config["outputs"])
                span.set_attributes({
                    "model": "test-llm-v1",
                    "temperature": 0.7,
                    "max_tokens": 100,
                })
                trace_id = span.trace_id
                created_trace_ids.append(trace_id)

                for exp_name, exp_value in trace_config["expectations"].items():
                    mlflow.log_expectation(
                        trace_id=trace_id,
                        name=exp_name,
                        value=exp_value,
                        span_id=span.span_id,
                        metadata={"trace_name": trace_config["name"]},
                    )

    dataset = create_dataset(
        name="comprehensive_trace_test",
        experiment_id=[experiment],
        tags={
            "test_type": "end_to_end",
            "model": "test-llm-v1",
            "language": "multilingual",
        },
    )

    traces = [client.get_trace(tid) for tid in created_trace_ids]
    dataset.merge_records(traces)

    df = dataset.to_df()
    assert len(df) == 3

    french_record = df[df["inputs"].apply(lambda x: "France" in str(x.get("question", "")))].iloc[0]
    assert french_record["expectations"]["correctness"] is True
    assert french_record["expectations"]["confidence_threshold"] == 0.8

    math_record = df[df["inputs"].apply(lambda x: "2+2" in str(x.get("question", "")))].iloc[0]
    assert math_record["expectations"]["correctness"] is False

    spanish_record = df[df["inputs"].apply(lambda x: x.get("language") == "es")].iloc[0]
    assert spanish_record["expectations"]["language_match"] is False
    assert spanish_record["expectations"]["politeness"] is True

    retrieved_dataset = get_dataset(dataset_id=dataset.dataset_id)
    retrieved_df = retrieved_dataset.to_df()
    assert len(retrieved_df) == 3
    assert retrieved_dataset.tags["model"] == "test-llm-v1"

    additional_records = [
        {
            "inputs": {"question": "What is Python?", "language": "en"},
            "expectations": {"technical_accuracy": True},
            "tags": {"source": "manual_addition"},
        }
    ]
    retrieved_dataset.merge_records(additional_records)

    final_df = retrieved_dataset.to_df()
    assert len(final_df) == 4

    trace_records = final_df[final_df["source_type"] == "TRACE"]
    assert len(trace_records) == 3

    manual_records = final_df[final_df["source_type"] != "TRACE"]
    assert len(manual_records) == 1


def test_dataset_pagination_transparency_large_records(experiments):
    dataset = create_dataset(
        name="test_pagination_transparency",
        experiment_id=experiments[0],
        tags={"test": "large_dataset"},
    )

    large_records = [
        {
            "inputs": {"question": f"Question {i}", "index": i},
            "expectations": {"answer": f"Answer {i}", "score": i * 0.01},
        }
        for i in range(150)
    ]

    dataset.merge_records(large_records)

    all_records = dataset._mlflow_dataset.records
    assert len(all_records) == 150

    record_indices = {record.inputs["index"] for record in all_records}
    expected_indices = set(range(150))
    assert record_indices == expected_indices

    record_scores = {record.expectations["score"] for record in all_records}
    expected_scores = {i * 0.01 for i in range(150)}
    assert record_scores == expected_scores

    df = dataset.to_df()
    assert len(df) == 150

    df_indices = {row["index"] for row in df["inputs"]}
    assert df_indices == expected_indices

    assert not hasattr(dataset, "page_token")
    assert not hasattr(dataset, "next_page_token")
    assert not hasattr(dataset, "max_results")

    second_access = dataset._mlflow_dataset.records
    assert second_access is all_records

    dataset._mlflow_dataset._records = None
    refreshed_records = dataset._mlflow_dataset.records
    assert len(refreshed_records) == 150


def test_dataset_internal_pagination_with_mock(experiments):
    from mlflow.tracking._tracking_service.utils import _get_store

    dataset = create_dataset(
        name="test_internal_pagination",
        experiment_id=experiments[0],
        tags={"test": "pagination_mock"},
    )

    records = [
        {"inputs": {"question": f"Q{i}", "id": i}, "expectations": {"answer": f"A{i}"}}
        for i in range(75)
    ]

    dataset.merge_records(records)

    dataset._mlflow_dataset._records = None

    store = _get_store()
    with mock.patch.object(
        store, "_load_dataset_records", wraps=store._load_dataset_records
    ) as mock_load:
        accessed_records = dataset._mlflow_dataset.records

        mock_load.assert_called_once_with(dataset.dataset_id, max_results=None)
        assert len(accessed_records) == 75

    dataset._mlflow_dataset._records = None

    with mock.patch.object(
        store, "_load_dataset_records", wraps=store._load_dataset_records
    ) as mock_load:
        df = dataset.to_df()

        mock_load.assert_called_once_with(dataset.dataset_id, max_results=None)
        assert len(df) == 75


def test_dataset_experiment_associations(experiments):
    from mlflow.genai.datasets import (
        add_dataset_to_experiments,
        remove_dataset_from_experiments,
    )

    dataset = create_dataset(
        name="test_associations",
        experiment_id=experiments[0],
        tags={"test": "associations"},
    )

    initial_exp_ids = dataset.experiment_ids
    assert experiments[0] in initial_exp_ids

    updated = add_dataset_to_experiments(
        dataset_id=dataset.dataset_id, experiment_ids=[experiments[1], experiments[2]]
    )
    assert experiments[0] in updated.experiment_ids
    assert experiments[1] in updated.experiment_ids
    assert experiments[2] in updated.experiment_ids
    assert len(updated.experiment_ids) == 3

    result = add_dataset_to_experiments(
        dataset_id=dataset.dataset_id, experiment_ids=[experiments[1], experiments[2]]
    )
    assert len(result.experiment_ids) == 3
    assert all(exp in result.experiment_ids for exp in experiments)

    removed = remove_dataset_from_experiments(
        dataset_id=dataset.dataset_id, experiment_ids=[experiments[1], experiments[2]]
    )
    assert experiments[1] not in removed.experiment_ids
    assert experiments[2] not in removed.experiment_ids
    assert experiments[0] in removed.experiment_ids
    assert len(removed.experiment_ids) == 1

    with mock.patch("mlflow.store.tracking.sqlalchemy_store._logger.warning") as mock_warning:
        idempotent = remove_dataset_from_experiments(
            dataset_id=dataset.dataset_id,
            experiment_ids=[experiments[1], experiments[2]],
        )
        assert mock_warning.call_count == 2
        assert "was not associated" in mock_warning.call_args_list[0][0][0]

    assert len(idempotent.experiment_ids) == 1


def test_dataset_associations_filestore_blocking(tmp_path):
    pytest.skip("FileStore is no longer supported.")
    from mlflow.genai.datasets import (
        add_dataset_to_experiments,
        remove_dataset_from_experiments,
    )

    mlflow.set_tracking_uri(tmp_path.as_uri())

    with pytest.raises(NotImplementedError, match="not supported with FileStore"):
        add_dataset_to_experiments(dataset_id="d-test123", experiment_ids=["1", "2"])

    with pytest.raises(NotImplementedError, match="not supported with FileStore"):
        remove_dataset_from_experiments(dataset_id="d-test123", experiment_ids=["1"])


def test_evaluation_dataset_tags_crud_workflow(experiments):
    dataset = create_dataset(
        name="test_tags_crud",
        experiment_id=experiments[0],
    )
    initial_tags = dataset.tags.copy()

    set_dataset_tags(
        dataset_id=dataset.dataset_id,
        tags={
            "team": "ml-platform",
            "project": "evaluation",
            "priority": "high",
        },
    )

    dataset = get_dataset(dataset_id=dataset.dataset_id)
    expected_tags = initial_tags.copy()
    expected_tags.update({
        "team": "ml-platform",
        "project": "evaluation",
        "priority": "high",
    })
    assert dataset.tags == expected_tags

    set_dataset_tags(
        dataset_id=dataset.dataset_id,
        tags={
            "priority": "medium",
            "status": "active",
        },
    )

    dataset = get_dataset(dataset_id=dataset.dataset_id)
    expected_tags = initial_tags.copy()
    expected_tags.update({
        "team": "ml-platform",
        "project": "evaluation",
        "priority": "medium",
        "status": "active",
    })
    assert dataset.tags == expected_tags

    delete_dataset_tag(
        dataset_id=dataset.dataset_id,
        key="priority",
    )

    dataset = get_dataset(dataset_id=dataset.dataset_id)
    expected_tags = initial_tags.copy()
    expected_tags.update({
        "team": "ml-platform",
        "project": "evaluation",
        "status": "active",
    })
    assert dataset.tags == expected_tags

    delete_dataset(dataset_id=dataset.dataset_id)

    with pytest.raises(MlflowException, match="Could not find|not found"):
        get_dataset(dataset_id=dataset.dataset_id)

    with pytest.raises(MlflowException, match="Could not find|not found"):
        set_dataset_tags(
            dataset_id=dataset.dataset_id,
            tags={"should": "fail"},
        )

    delete_dataset_tag(dataset_id=dataset.dataset_id, key="status")


def test_set_dataset_tags_databricks(mock_databricks_environment):
    with pytest.raises(NotImplementedError, match="tag operations are not available"):
        set_dataset_tags(dataset_id="test", tags={"key": "value"})


def test_delete_dataset_tag_databricks(mock_databricks_environment):
    with pytest.raises(NotImplementedError, match="tag operations are not available"):
        delete_dataset_tag(dataset_id="test", key="key")


def test_dataset_schema_evolution_and_log_input(experiments):
    dataset = create_dataset(
        name="schema_evolution_test",
        experiment_id=[experiments[0]],
        tags={"test": "schema_evolution", "mlflow.user": "test_user"},
    )

    stage1_records = [
        {
            "inputs": {"prompt": "What is MLflow?"},
            "expectations": {"response": "MLflow is a platform"},
        }
    ]
    dataset.merge_records(stage1_records)

    ds1 = get_dataset(dataset_id=dataset.dataset_id)
    schema1 = json.loads(ds1.schema)
    assert schema1 is not None
    assert "prompt" in schema1["inputs"]
    assert schema1["inputs"]["prompt"] == "string"
    assert len(schema1["inputs"]) == 1
    assert len(schema1["expectations"]) == 1

    stage2_records = [
        {
            "inputs": {
                "prompt": "Explain Python",
                "temperature": 0.7,
                "max_length": 500,
                "top_p": 0.95,
            },
            "expectations": {
                "response": "Python is a programming language",
                "quality_score": 0.85,
                "token_count": 127,
            },
        }
    ]
    dataset.merge_records(stage2_records)

    ds2 = get_dataset(dataset_id=dataset.dataset_id)
    schema2 = json.loads(ds2.schema)
    assert "temperature" in schema2["inputs"]
    assert schema2["inputs"]["temperature"] == "float"
    assert "max_length" in schema2["inputs"]
    assert schema2["inputs"]["max_length"] == "integer"
    assert len(schema2["inputs"]) == 4
    assert len(schema2["expectations"]) == 3

    stage3_records = [
        {
            "inputs": {
                "prompt": "Complex query",
                "streaming": True,
                "stop_sequences": ["\n\n", "END"],
                "config": {"model": "gpt-4", "version": "1.0"},
            },
            "expectations": {
                "response": "Complex response",
                "is_valid": True,
                "citations": ["source1", "source2"],
                "metadata": {"confidence": 0.9},
            },
        }
    ]
    dataset.merge_records(stage3_records)

    ds3 = get_dataset(dataset_id=dataset.dataset_id)
    schema3 = json.loads(ds3.schema)

    assert schema3["inputs"]["streaming"] == "boolean"
    assert schema3["inputs"]["stop_sequences"] == "array"
    assert schema3["inputs"]["config"] == "object"
    assert schema3["expectations"]["is_valid"] == "boolean"
    assert schema3["expectations"]["citations"] == "array"
    assert schema3["expectations"]["metadata"] == "object"

    assert "prompt" in schema3["inputs"]
    assert "temperature" in schema3["inputs"]
    assert "quality_score" in schema3["expectations"]

    with mlflow.start_run(experiment_id=experiments[0]) as run:
        mlflow.log_input(dataset, context="evaluation")

        mlflow.log_metrics({"accuracy": 0.92, "f1_score": 0.89})

    run_data = mlflow.get_run(run.info.run_id)
    assert run_data.inputs is not None
    assert run_data.inputs.dataset_inputs is not None
    assert len(run_data.inputs.dataset_inputs) > 0

    dataset_input = run_data.inputs.dataset_inputs[0]
    assert dataset_input.dataset.name == "schema_evolution_test"
    assert dataset_input.dataset.source_type == "mlflow_evaluation_dataset"

    tag_dict = {tag.key: tag.value for tag in dataset_input.tags}
    assert "mlflow.data.context" in tag_dict
    assert tag_dict["mlflow.data.context"] == "evaluation"

    final_dataset = get_dataset(dataset_id=dataset.dataset_id)
    final_schema = json.loads(final_dataset.schema)

    assert "inputs" in final_schema
    assert "expectations" in final_schema
    assert "version" in final_schema
    assert final_schema["version"] == "1.0"

    profile = json.loads(final_dataset.profile)
    assert profile is not None
    assert profile["num_records"] == 3

    consistency_records = [
        {
            "inputs": {"prompt": "Another test", "temperature": 0.5, "max_length": 200},
            "expectations": {"response": "Another response", "quality_score": 0.75},
        }
    ]
    dataset.merge_records(consistency_records)

    consistent_dataset = get_dataset(dataset_id=dataset.dataset_id)
    consistent_schema = json.loads(consistent_dataset.schema)

    assert set(consistent_schema["inputs"].keys()) == set(final_schema["inputs"].keys())
    assert set(consistent_schema["expectations"].keys()) == set(final_schema["expectations"].keys())

    consistent_profile = json.loads(consistent_dataset.profile)
    assert consistent_profile["num_records"] == 4
    delete_dataset_tag(dataset_id="test", key="key")


def test_deprecated_parameter_substitution(experiment):
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")

        dataset = create_dataset(
            uc_table_name="test_dataset_deprecated",
            experiment_id=experiment,
            tags={"test": "deprecated_parameter"},
        )

        assert len(w) == 1
        assert issubclass(w[0].category, FutureWarning)
        assert "uc_table_name" in str(w[0].message)
        assert "deprecated" in str(w[0].message).lower()
        assert "name" in str(w[0].message)

        assert dataset.name == "test_dataset_deprecated"
        assert dataset.tags["test"] == "deprecated_parameter"

    with pytest.raises(ValueError, match="Cannot specify both.*uc_table_name.*and.*name"):
        create_dataset(
            uc_table_name="old_name",
            name="new_name",
            experiment_id=experiment,
        )

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")

        with pytest.raises(ValueError, match="name.*only supported in Databricks"):
            delete_dataset(uc_table_name="test_dataset_deprecated")

        assert len(w) == 1
        assert issubclass(w[0].category, FutureWarning)
        assert "uc_table_name" in str(w[0].message)

    delete_dataset(dataset_id=dataset.dataset_id)


def test_create_dataset_uses_active_experiment_when_not_specified(client):
    exp_id = mlflow.create_experiment("test_active_experiment")
    mlflow.set_experiment(experiment_id=exp_id)

    dataset = create_dataset(name="test_with_active_exp")

    assert dataset.experiment_ids == [exp_id]

    from mlflow.tracking import fluent

    fluent._active_experiment_id = None


def test_create_dataset_with_no_active_experiment(client):
    from mlflow.tracking import fluent

    fluent._active_experiment_id = None

    dataset = create_dataset(name="test_no_active_exp")

    assert dataset.experiment_ids == ["0"]


def test_create_dataset_explicit_overrides_active_experiment(client):
    active_exp = mlflow.create_experiment("active_exp")
    explicit_exp = mlflow.create_experiment("explicit_exp")

    mlflow.set_experiment(experiment_id=active_exp)

    dataset = create_dataset(name="test_explicit_override", experiment_id=explicit_exp)

    assert dataset.experiment_ids == [explicit_exp]

    from mlflow.tracking import fluent

    fluent._active_experiment_id = None


def test_create_dataset_none_uses_active_experiment(client):
    exp_id = mlflow.create_experiment("test_none_experiment")
    mlflow.set_experiment(experiment_id=exp_id)

    dataset = create_dataset(name="test_none_exp", experiment_id=None)

    assert dataset.experiment_ids == [exp_id]

    from mlflow.tracking import fluent

    fluent._active_experiment_id = None


def test_source_type_inference():
    exp = mlflow.create_experiment("test_source_inference")
    dataset = create_dataset(
        name="test_source_inference",
        experiment_id=exp,
        tags={"test": "source_inference"},
    )

    human_records = [
        {
            "inputs": {"question": "What is MLflow?"},
            "expectations": {"answer": "MLflow is an ML platform", "quality": 0.9},
        },
        {
            "inputs": {"question": "How to track experiments?"},
            "expectations": {"answer": "Use mlflow.start_run()", "quality": 0.85},
        },
    ]
    dataset.merge_records(human_records)

    df = dataset.to_df()
    human_sources = df[df["source_type"] == DatasetRecordSourceType.HUMAN.value]
    assert len(human_sources) == 2

    code_records = [{"inputs": {"question": f"Generated question {i}"}} for i in range(3)]
    dataset.merge_records(code_records)

    df = dataset.to_df()
    code_sources = df[df["source_type"] == DatasetRecordSourceType.CODE.value]
    assert len(code_sources) == 3

    explicit_records = [
        {
            "inputs": {"question": "Document-based question"},
            "expectations": {"answer": "From document"},
            "source": {
                "source_type": DatasetRecordSourceType.DOCUMENT.value,
                "source_data": {"source_id": "doc123", "page": 5},
            },
        }
    ]
    dataset.merge_records(explicit_records)

    df = dataset.to_df()
    doc_sources = df[df["source_type"] == DatasetRecordSourceType.DOCUMENT.value]
    assert len(doc_sources) == 1
    assert doc_sources.iloc[0]["source_id"] == "doc123"

    empty_exp_records = [{"inputs": {"question": "Has empty expectations"}, "expectations": {}}]
    dataset.merge_records(empty_exp_records)

    df = dataset.to_df()
    last_record = df.iloc[-1]
    assert last_record["source_type"] not in [
        DatasetRecordSourceType.HUMAN.value,
        DatasetRecordSourceType.CODE.value,
    ]

    explicit_trace = [
        {
            "inputs": {"question": "From trace"},
            "source": {
                "source_type": DatasetRecordSourceType.TRACE.value,
                "source_data": {"trace_id": "trace123"},
            },
        }
    ]
    dataset.merge_records(explicit_trace)

    df = dataset.to_df()
    trace_sources = df[df["source_type"] == DatasetRecordSourceType.TRACE.value]
    assert len(trace_sources) == 1, f"Expected 1 TRACE source, got {len(trace_sources)}"
    assert trace_sources.iloc[0]["source_id"] == "trace123"

    source_counts = df["source_type"].value_counts()
    assert source_counts.get(DatasetRecordSourceType.HUMAN.value, 0) == 2
    assert source_counts.get(DatasetRecordSourceType.CODE.value, 0) == 3
    assert source_counts.get(DatasetRecordSourceType.DOCUMENT.value, 0) == 1
    assert source_counts.get(DatasetRecordSourceType.TRACE.value, 0) == 1

    delete_dataset(dataset_id=dataset.dataset_id)


def test_trace_source_type_detection():
    exp = mlflow.create_experiment("test_trace_source_detection")

    trace_ids = []
    for i in range(3):
        with mlflow.start_run(experiment_id=exp):
            with mlflow.start_span(name=f"test_span_{i}") as span:
                span.set_inputs({"question": f"Question {i}", "context": f"Context {i}"})
                span.set_outputs({"answer": f"Answer {i}"})
                trace_ids.append(span.trace_id)

                if i < 2:
                    mlflow.log_expectation(
                        trace_id=span.trace_id,
                        name="quality",
                        value=0.8 + i * 0.05,
                        span_id=span.span_id,
                    )

    dataset = create_dataset(
        name="test_trace_sources",
        experiment_id=exp,
        tags={"test": "trace_source_detection"},
    )

    client = mlflow.MlflowClient()
    traces = [client.get_trace(tid) for tid in trace_ids]
    dataset.merge_records(traces)

    df = dataset.to_df()
    trace_sources = df[df["source_type"] == DatasetRecordSourceType.TRACE.value]
    assert len(trace_sources) == 3

    for trace_id in trace_ids:
        matching_records = df[df["source_id"] == trace_id]
        assert len(matching_records) == 1

    dataset2 = create_dataset(
        name="test_trace_sources_df",
        experiment_id=exp,
        tags={"test": "trace_source_df"},
    )

    traces_df = mlflow.search_traces(locations=[exp])
    assert not traces_df.empty
    dataset2.merge_records(traces_df)

    df2 = dataset2.to_df()
    trace_sources2 = df2[df2["source_type"] == DatasetRecordSourceType.TRACE.value]
    assert len(trace_sources2) == len(traces_df)

    dataset3 = create_dataset(
        name="test_trace_sources_list",
        experiment_id=exp,
        tags={"test": "trace_source_list"},
    )

    traces_list = mlflow.search_traces(locations=[exp], return_type="list")
    assert len(traces_list) > 0
    dataset3.merge_records(traces_list)

    df3 = dataset3.to_df()
    trace_sources3 = df3[df3["source_type"] == DatasetRecordSourceType.TRACE.value]
    assert len(trace_sources3) == len(traces_list)

    df_with_expectations = df[df["expectations"].apply(lambda x: bool(x) and len(x) > 0)]
    assert len(df_with_expectations) == 2

    delete_dataset(dataset_id=dataset.dataset_id)
    delete_dataset(dataset_id=dataset2.dataset_id)
    delete_dataset(dataset_id=dataset3.dataset_id)


def test_create_dataset_empty_list_stays_empty(client):
    exp_id = mlflow.create_experiment("test_empty_list")
    mlflow.set_experiment(experiment_id=exp_id)

    dataset = create_dataset(name="test_empty_list", experiment_id=[])

    assert dataset.experiment_ids == []

    from mlflow.tracking import fluent

    fluent._active_experiment_id = None


def test_search_datasets_filter_string_edge_cases(client):
    exp_id = mlflow.create_experiment("test_filter_edge_cases")

    dataset = create_dataset(name="test_dataset", experiment_id=exp_id, tags={"test": "value"})

    with mock.patch("mlflow.tracking.client.MlflowClient.search_datasets") as mock_search:
        mock_search.return_value = mock.MagicMock(token=None, items=[dataset])

        search_datasets(experiment_ids=exp_id, filter_string=None)
        call_args = mock_search.call_args
        filter_arg = call_args.kwargs.get("filter_string")
        assert "created_time >=" in filter_arg

        mock_search.reset_mock()

        search_datasets(experiment_ids=exp_id, filter_string=[])
        call_args = mock_search.call_args
        filter_arg = call_args.kwargs.get("filter_string")
        assert "created_time >=" in filter_arg

        mock_search.reset_mock()

        search_datasets(experiment_ids=exp_id, filter_string="")
        call_args = mock_search.call_args
        filter_arg = call_args.kwargs.get("filter_string")
        assert "created_time >=" in filter_arg

        mock_search.reset_mock()

        search_datasets(experiment_ids=exp_id, filter_string='name = "test"')
        call_args = mock_search.call_args
        filter_arg = call_args.kwargs.get("filter_string")
        assert filter_arg == 'name = "test"'


def test_wrapper_type_is_actually_returned_not_entity(experiments):
    dataset = create_dataset(
        name="test_wrapper",
        experiment_id=experiments[0],
        tags={"test": "wrapper_check"},
    )

    assert isinstance(dataset, WrapperEvaluationDataset)
    assert not isinstance(dataset, EntityEvaluationDataset)
    assert hasattr(dataset, "_mlflow_dataset")
    assert dataset._mlflow_dataset is not None
    assert isinstance(dataset._mlflow_dataset, EntityEvaluationDataset)


def test_wrapper_delegates_all_properties_correctly(experiments):
    dataset = create_dataset(
        name="test_delegation",
        experiment_id=experiments[0],
        tags={"env": "test", "version": "1.0"},
    )

    assert dataset.name == "test_delegation"
    assert dataset.dataset_id.startswith("d-")
    assert dataset.tags["env"] == "test"
    assert dataset.tags["version"] == "1.0"
    assert experiments[0] in dataset.experiment_ids
    assert dataset.created_time > 0
    assert dataset.last_update_time > 0
    assert dataset.digest is not None
    assert hasattr(dataset, "source")
    assert dataset.source._get_source_type() == "mlflow_evaluation_dataset"


def test_get_and_search_return_wrapper_not_entity(experiments):
    created = create_dataset(
        name="test_get_wrapper",
        experiment_id=experiments[0],
        tags={"test": "get"},
    )

    retrieved = get_dataset(dataset_id=created.dataset_id)
    assert isinstance(retrieved, WrapperEvaluationDataset)
    assert not isinstance(retrieved, EntityEvaluationDataset)
    assert retrieved.dataset_id == created.dataset_id
    assert retrieved.name == created.name

    results = search_datasets(
        experiment_ids=experiments[0],
        filter_string="name = 'test_get_wrapper'",
    )
    assert len(results) == 1
    assert isinstance(results[0], WrapperEvaluationDataset)
    assert not isinstance(results[0], EntityEvaluationDataset)


def test_wrapper_vs_direct_client_usage(experiments):
    client = MlflowClient()

    entity_dataset = client.create_dataset(
        name="test_client_direct",
        experiment_id=experiments[0],
        tags={"direct": "client"},
    )
    assert isinstance(entity_dataset, EntityEvaluationDataset)
    assert not isinstance(entity_dataset, WrapperEvaluationDataset)

    wrapped_dataset = create_dataset(
        name="test_wrapped",
        experiment_id=experiments[0],
        tags={"wrapped": "fluent"},
    )
    assert isinstance(wrapped_dataset, WrapperEvaluationDataset)
    assert not isinstance(wrapped_dataset, EntityEvaluationDataset)
    assert wrapped_dataset._mlflow_dataset is not None

    wrapped_from_entity = WrapperEvaluationDataset(entity_dataset)
    assert wrapped_from_entity == entity_dataset


def test_wrapper_works_with_mlflow_log_input_integration(experiments):
    dataset = create_dataset(
        name="test_log_input",
        experiment_id=experiments[0],
    )

    records = [
        {
            "inputs": {"question": "Test question"},
            "expectations": {"answer": "Test answer"},
        }
    ]
    dataset.merge_records(records)

    with mlflow.start_run(experiment_id=experiments[0]) as run:
        mlflow.log_input(dataset, context="evaluation")

    run_data = mlflow.get_run(run.info.run_id)
    assert len(run_data.inputs.dataset_inputs) == 1
    dataset_input = run_data.inputs.dataset_inputs[0]
    assert dataset_input.dataset.name == "test_log_input"
    assert dataset_input.dataset.digest == dataset.digest


def test_wrapper_isinstance_checks_for_dataset_interfaces(experiments):
    dataset = create_dataset(
        name="test_isinstance",
        experiment_id=experiments[0],
    )

    assert isinstance(dataset, Dataset)
    assert isinstance(dataset, PyFuncConvertibleDatasetMixin)
    assert isinstance(dataset, WrapperEvaluationDataset)
    assert not isinstance(dataset, EntityEvaluationDataset)
    assert isinstance(dataset, (WrapperEvaluationDataset, EntityEvaluationDataset))


@pytest.mark.parametrize(
    "records",
    [
        [
            {"inputs": {"persona": "Student", "goal": "Find articles"}},
            {
                "inputs": {
                    "persona": "Researcher",
                    "goal": "Review",
                    "context": {"dept": "CS"},
                }
            },
            {"inputs": {"goal": "Single goal"}, "expectations": {"output": "expected"}},
        ],
        [
            {"inputs": {"goal": "Learn ML", "simulation_guidelines": "Be concise"}},
            {
                "inputs": {
                    "persona": "Engineer",
                    "goal": "Debug",
                    "simulation_guidelines": "Focus on logs",
                }
            },
            {
                "inputs": {
                    "persona": "Student",
                    "goal": "Study",
                    "context": {"course": "CS101"},
                    "simulation_guidelines": "Ask clarifying questions",
                }
            },
        ],
    ],
)
def test_multiturn_valid_formats(experiments, records):
    dataset = create_dataset(name="multiturn_test", experiment_id=experiments[0])
    dataset.merge_records(records)
    df = dataset.to_df()

    assert len(df) == 3
    for _, row in df.iterrows():
        assert any(
            key in row["inputs"] for key in ["persona", "goal", "context", "simulation_guidelines"]
        )


@pytest.mark.parametrize(
    ("records", "error_pattern"),
    [
        # Top-level session fields
        (
            [{"persona": "Student", "goal": "Find articles", "custom_field": "value"}],
            "Each record must have an 'inputs' field",
        ),
        # Mixed fields in inputs
        (
            [
                {
                    "inputs": {
                        "persona": "Student",
                        "goal": "Find",
                        "custom_field": "value",
                    }
                }
            ],
            "Invalid input schema.*cannot mix session fields",
        ),
        # Inconsistent batch schema
        (
            [
                {"inputs": {"persona": "Student", "goal": "Find articles"}},
                {"inputs": {"question": "What is MLflow?"}},
            ],
            "must use the same granularity.*Found",
        ),
        # Empty inputs in batch with session records
        (
            [
                {"inputs": {"goal": "Find articles"}},
                {"inputs": {}},
            ],
            "Empty inputs are not allowed for session records.*'goal' field is required",
        ),
    ],
)
def test_multiturn_validation_errors(experiments, records, error_pattern):
    dataset = create_dataset(name="multiturn_error_test", experiment_id=experiments[0])
    with pytest.raises(MlflowException, match=error_pattern):
        dataset.merge_records(records)


@pytest.mark.parametrize(
    ("existing_records", "new_records"),
    [
        # Multiturn then custom
        (
            [{"inputs": {"persona": "Student", "goal": "Find articles"}}],
            [{"inputs": {"question": "What is MLflow?", "model": "gpt-4"}}],
        ),
        # Custom then multiturn
        (
            [{"inputs": {"question": "What is MLflow?", "model": "gpt-4"}}],
            [{"inputs": {"persona": "Student", "goal": "Find articles"}}],
        ),
    ],
)
def test_multiturn_schema_compatibility(experiments, existing_records, new_records):
    dataset = create_dataset(name="multiturn_compat_test", experiment_id=experiments[0])
    dataset.merge_records(existing_records)

    with pytest.raises(MlflowException, match="Cannot mix granularities"):
        dataset.merge_records(new_records)


def test_multiturn_with_expectations_and_tags(experiments):
    dataset = create_dataset(name="multiturn_full_test", experiment_id=experiments[0])
    records = [
        {
            "inputs": {
                "persona": "Graduate Student",
                "goal": "Find peer-reviewed articles on machine learning",
                "context": {"user_id": "U0001", "department": "CS"},
                "simulation_guidelines": "Be thorough and cite sources",
            },
            "expectations": {"expected_output": "relevant articles", "quality": "high"},
            "tags": {"difficulty": "medium"},
        },
        {
            "inputs": {
                "persona": "Librarian",
                "goal": "Help with inter-library loan",
            },
            "expectations": {"expected_output": "loan information"},
        },
    ]

    dataset.merge_records(records)

    df = dataset.to_df()
    assert len(df) == 2

    grad_record = df[df["inputs"].apply(lambda x: x.get("persona") == "Graduate Student")].iloc[0]
    assert grad_record["expectations"]["expected_output"] == "relevant articles"
    assert grad_record["expectations"]["quality"] == "high"
    assert grad_record["tags"]["difficulty"] == "medium"
    assert grad_record["inputs"]["context"] == {"user_id": "U0001", "department": "CS"}
    assert grad_record["inputs"]["simulation_guidelines"] == "Be thorough and cite sources"