Source code for mlflow.genai.datasets

"""
Databricks Agent Datasets Python SDK. For more details see Databricks Agent Evaluation:
 <https://docs.databricks.com/en/generative-ai/agent-evaluation/index.html>

The API docs can be found here:
<https://api-docs.databricks.com/python/databricks-agents/latest/databricks_agent_eval.html#datasets>
"""

from typing import Optional, Union

from mlflow.genai.datasets.evaluation_dataset import EvaluationDataset

_ERROR_MSG = (
    "The `databricks-agents` package is required to use `mlflow.genai.datasets`. "
    "Please install it with `pip install databricks-agents`."
)


[docs]def create_dataset( uc_table_name: str, experiment_id: Optional[Union[str, list[str]]] = None ) -> "EvaluationDataset": """Create a dataset with the given name and associate it with the given experiment. Args: uc_table_name: The UC table name of the dataset. experiment_id: The ID of the experiment to associate the dataset with. If not provided, the current experiment is inferred from the environment. """ try: from databricks.agents.datasets import create_dataset except ImportError as e: raise ImportError(_ERROR_MSG) from e return EvaluationDataset(create_dataset(uc_table_name, experiment_id))
[docs]def delete_dataset(uc_table_name: str) -> None: """Delete the dataset with the given name. Args: uc_table_name: The UC table name of the dataset. """ try: from databricks.agents.datasets import delete_dataset except ImportError: raise ImportError(_ERROR_MSG) from None return delete_dataset(uc_table_name)
[docs]def get_dataset(uc_table_name: str) -> "EvaluationDataset": """Get the dataset with the given name. Args: uc_table_name: The UC table name of the dataset. """ try: from databricks.agents.datasets import get_dataset except ImportError as e: raise ImportError(_ERROR_MSG) from e return EvaluationDataset(get_dataset(uc_table_name))
__all__ = [ "create_dataset", "delete_dataset", "get_dataset", "EvaluationDataset", ]