Dataset Management

def create_dataset(
    self,
    dataset_name: str,
    *,
    description: Optional[str] = None,
    data_type: Optional[DataType] = None,
    inputs_schema: Optional[dict] = None,
    outputs_schema: Optional[dict] = None,
    metadata: Optional[dict] = None,
    tags: Optional[list[str]] = None,
) -> Dataset:
    """
    Create a new dataset.

    Parameters:
    - dataset_name: Name of the dataset
    - description: Dataset description
    - data_type: Type of data (e.g., "kv", "llm", "chat")
    - inputs_schema: JSON schema for inputs
    - outputs_schema: JSON schema for outputs
    - metadata: Dataset metadata
    - tags: Dataset tags

    Returns:
    Created Dataset object
    """

def read_dataset(
    self,
    dataset_name: Optional[str] = None,
    dataset_id: Optional[Union[str, UUID]] = None
) -> Dataset:
    """
    Read a dataset by name or ID.

    Parameters:
    - dataset_name: Name of the dataset
    - dataset_id: ID of the dataset (alternative to name)

    Returns:
    Dataset object
    """

def list_datasets(
    self,
    *,
    dataset_ids: Optional[list[Union[str, UUID]]] = None,
    dataset_name: Optional[str] = None,
    dataset_name_contains: Optional[str] = None,
    data_type: Optional[str] = None,
    limit: Optional[int] = None,
    offset: Optional[int] = None,
    tags: Optional[list[str]] = None,
) -> Iterator[Dataset]:
    """
    List all datasets with optional filtering.

    Parameters:
    - dataset_ids: Filter by specific dataset IDs
    - dataset_name: Filter by exact name
    - dataset_name_contains: Filter by name substring
    - data_type: Filter by data type
    - limit: Maximum number of datasets
    - offset: Number of datasets to skip
    - tags: Filter by tags

    Returns:
    Iterator of Dataset objects
    """

def delete_dataset(
    self,
    dataset_id: Optional[Union[str, UUID]] = None,
    dataset_name: Optional[str] = None
) -> None:
    """
    Delete a dataset.

    Parameters:
    - dataset_id: ID of the dataset
    - dataset_name: Name of the dataset (alternative to ID)
    """

def has_dataset(
    self,
    dataset_name: str
) -> bool:
    """
    Check if a dataset exists.

    Parameters:
    - dataset_name: Name of the dataset

    Returns:
    True if dataset exists
    """

def update_dataset_tag(
    self,
    dataset_id: Union[str, UUID],
    tag: str,
    *,
    as_of: Optional[Union[datetime, str]] = None,
) -> None:
    """
    Update dataset version tag.

    Parameters:
    - dataset_id: ID of the dataset
    - tag: Tag name for this version
    - as_of: Timestamp for this version
    """

def list_dataset_versions(
    self,
    dataset_id: Union[str, UUID],
    *,
    limit: Optional[int] = None,
    offset: Optional[int] = None,
    search: Optional[str] = None,
) -> Iterator[dict]:
    """
    List versions of a dataset.

    Parameters:
    - dataset_id: ID of the dataset
    - limit: Maximum number of versions
    - offset: Number of versions to skip
    - search: Search string for tags

    Returns:
    Iterator of version info dictionaries
    """

def read_dataset_version(
    self,
    dataset_id: Union[str, UUID],
    version: Union[str, datetime],
    *,
    include_examples: bool = False,
) -> dict:
    """
    Read a specific dataset version.

    Parameters:
    - dataset_id: ID of the dataset
    - version: Version tag or timestamp
    - include_examples: Whether to include examples

    Returns:
    Version info dictionary
    """

def diff_dataset_versions(
    self,
    dataset_id: Union[str, UUID],
    from_version: Union[str, datetime],
    to_version: Union[str, datetime],
) -> dict:
    """
    Diff two dataset versions.

    Parameters:
    - dataset_id: ID of the dataset
    - from_version: Starting version
    - to_version: Ending version

    Returns:
    Diff information dictionary
    """

def clone_public_dataset(
    self,
    token: str,
    *,
    source_api_url: Optional[str] = None,
    dataset_name: Optional[str] = None,
) -> Dataset:
    """
    Clone a public dataset.

    Parameters:
    - token: Public dataset token
    - source_api_url: Source API URL
    - dataset_name: Custom name for cloned dataset

    Returns:
    Cloned Dataset object
    """

def upload_dataframe(
    self,
    df: Any,
    *,
    name: str,
    description: Optional[str] = None,
    input_keys: Sequence[str],
    output_keys: Sequence[str],
    data_type: Optional[DataType] = None,
) -> Dataset:
    """
    Upload a pandas DataFrame as a dataset.

    Parameters:
    - df: pandas DataFrame
    - name: Dataset name
    - description: Dataset description
    - input_keys: Column names to use as inputs
    - output_keys: Column names to use as outputs
    - data_type: Type of data

    Returns:
    Created Dataset object
    """

def upload_csv(
    self,
    csv_file: Union[str, Path],
    *,
    name: str,
    description: Optional[str] = None,
    input_keys: Sequence[str],
    output_keys: Sequence[str],
    data_type: Optional[DataType] = None,
) -> Dataset:
    """
    Upload a CSV file as a dataset.

    Parameters:
    - csv_file: Path to CSV file
    - name: Dataset name
    - description: Dataset description
    - input_keys: Column names to use as inputs
    - output_keys: Column names to use as outputs
    - data_type: Type of data

    Returns:
    Created Dataset object
    """