arize 8.0.0a22__py3-none-any.whl → 8.0.0b0__py3-none-any.whl

arize/datasets/client.py

@@ -1,18 +1,20 @@
+"""Client implementation for managing datasets in the Arize platform."""
+
 from __future__ import annotations
 
 import logging
 import time
 import uuid
-from typing import Any, Dict, List
+from typing import TYPE_CHECKING
 
 import pandas as pd
 import pyarrow as pa
 
 from arize._flight.client import ArizeFlightClient
 from arize._generated.api_client import models
-from arize.config import SDKConfiguration
 from arize.datasets.validation import validate_dataset_df
 from arize.exceptions.base import INVALID_ARROW_CONVERSION_MSG
+from arize.pre_releases import ReleaseStage, prerelease_endpoint
 from arize.utils.cache import cache_resource, load_cached_resource
 from arize.utils.openinference_conversion import (
     convert_boolean_columns_to_str,
@@ -21,40 +23,121 @@ from arize.utils.openinference_conversion import (
 )
 from arize.utils.size import get_payload_size_mb
 
+if TYPE_CHECKING:
+    from arize.config import SDKConfiguration
+
 logger = logging.getLogger(__name__)
 
 
 class DatasetsClient:
-    def __init__(self, *, sdk_config: SDKConfiguration):
+    """Client for managing datasets including creation, retrieval, and example management.
+
+    This class is primarily intended for internal use within the SDK. Users are
+    highly encouraged to access resource-specific functionality via
+    :class:`arize.ArizeClient`.
+
+    The datasets client is a thin wrapper around the generated REST API client,
+    using the shared generated API client owned by
+    :class:`arize.config.SDKConfiguration`.
+    """
+
+    def __init__(self, *, sdk_config: SDKConfiguration) -> None:
+        """
+        Args:
+            sdk_config: Resolved SDK configuration.
+        """  # noqa: D205, D212
         self._sdk_config = sdk_config
 
-        # Import at runtime so it’s still lazy and extras-gated by the parent
+        # Import at runtime so it's still lazy and extras-gated by the parent
         from arize._generated import api_client as gen
 
         # Use the shared generated client from the config
         self._api = gen.DatasetsApi(self._sdk_config.get_generated_client())
 
-        # Forward methods to preserve exact runtime signatures/docs
-        self.list = self._api.datasets_list
-        self.get = self._api.datasets_get
-        self.delete = self._api.datasets_delete
+    @prerelease_endpoint(key="datasets.list", stage=ReleaseStage.BETA)
+    def list(
+        self,
+        *,
+        space_id: str | None = None,
+        limit: int = 100,
+        cursor: str | None = None,
+    ) -> models.DatasetsList200Response:
+        """List datasets the user has access to.
+
+        Datasets are returned in descending creation order (most recently created
+        first). Dataset versions are not included in this response; use `get()` to
+        retrieve a dataset along with its versions.
+
+        Args:
+            space_id: Optional space ID to scope results to a single space.
+            limit: Maximum number of datasets to return. The server enforces an
+                upper bound.
+            cursor: Opaque pagination cursor returned from a previous response.
 
-        # Custom methods
-        self.create = self._create_dataset
-        self.list_examples = self._list_examples
+        Returns:
+            A response object with the datasets and pagination information.
 
-    def _create_dataset(
+        Raises:
+            arize._generated.api_client.exceptions.ApiException: If the REST API
+                returns an error response (e.g. 401/403/429).
+        """
+        return self._api.datasets_list(
+            space_id=space_id,
+            limit=limit,
+            cursor=cursor,
+        )
+
+    @prerelease_endpoint(key="datasets.create", stage=ReleaseStage.BETA)
+    def create(
         self,
         *,
         name: str,
         space_id: str,
-        examples: List[Dict[str, Any]] | pd.DataFrame,
+        examples: list[dict[str, object]] | pd.DataFrame,
         force_http: bool = False,
-    ):
-        if not isinstance(examples, (list, pd.DataFrame)):
+    ) -> models.Dataset:
+        """Create a dataset with JSON examples.
+
+        Empty datasets are not allowed.
+
+        Payload notes (server-enforced):
+            - `name` must be unique within the given `space_id`.
+            - Each example may contain arbitrary user-defined fields.
+            - Do not include system-managed fields on create: `id`, `created_at`,
+              `updated_at` (requests containing these fields will be rejected).
+            - Each example must contain at least one property (i.e. `{}` is invalid).
+
+        Transport selection:
+            - If the payload is below the configured REST payload threshold (or
+              `force_http=True`), this method uploads via REST.
+            - Otherwise, it attempts a more efficient upload path via gRPC + Flight.
+
+        Args:
+            name: Dataset name (must be unique within the target space).
+            space_id: Space ID to create the dataset in.
+            examples: Dataset examples either as:
+                - a list of JSON-like dicts, or
+                - a pandas DataFrame (will be converted to records for REST).
+            force_http: If True, force REST upload even if the payload exceeds the
+                configured REST payload threshold.
+
+        Returns:
+            The created dataset object as returned by the API.
+
+        Raises:
+            TypeError: If `examples` is not a list of dicts or a pandas DataFrame.
+            RuntimeError: If the Flight upload path is selected and the Flight request
+                fails.
+            arize._generated.api_client.exceptions.ApiException: If the REST API
+                returns an error response (e.g. 400/401/403/409/429).
+        """
+        if not isinstance(examples, list | pd.DataFrame):
             raise TypeError(
                 "Examples must be a list of dicts or a pandas DataFrame"
             )
+        if len(examples) == 0:
+            raise ValueError("Cannot create an empty dataset")
+
         below_threshold = (
             get_payload_size_mb(examples)
             <= self._sdk_config.max_http_payload_size_mb
@@ -63,14 +146,14 @@ class DatasetsClient:
             from arize._generated import api_client as gen
 
             data = (
-                examples.to_dict(orient="records")
+                examples.to_dict(orient="records")  # type: ignore
                 if isinstance(examples, pd.DataFrame)
                 else examples
             )
 
             body = gen.DatasetsCreateRequest(
                 name=name,
-                spaceId=space_id,
+                space_id=space_id,
                 examples=data,
             )
             return self._api.datasets_create(datasets_create_request=body)
@@ -93,76 +176,83 @@ class DatasetsClient:
             examples=data,
         )
 
-    def _create_dataset_via_flight(
-        self,
-        name: str,
-        space_id: str,
-        examples: pd.DataFrame,
-    ):
-        data = examples.copy()
-        # Convert datetime columns to int64 (ms since epoch)
-        data = convert_datetime_columns_to_int(data)
-        data = convert_boolean_columns_to_str(data)
-        data = _set_default_columns_for_dataset(data)
-        data = convert_default_columns_to_json_str(data)
+    @prerelease_endpoint(key="datasets.get", stage=ReleaseStage.BETA)
+    def get(self, *, dataset_id: str) -> models.Dataset:
+        """Get a dataset by ID.
 
-        validation_errors = validate_dataset_df(data)
-        if validation_errors:
-            raise RuntimeError([e.error_message() for e in validation_errors])
+        The returned dataset includes its dataset versions (sorted by creation time,
+        most recent first). Dataset examples are not included; use `list_examples()`
+        to retrieve examples.
 
-        # Convert to Arrow table
-        try:
-            logger.debug("Converting data to Arrow format")
-            pa_table = pa.Table.from_pandas(data, preserve_index=False)
-        except pa.ArrowInvalid as e:
-            logger.error(f"{INVALID_ARROW_CONVERSION_MSG}: {str(e)}")
-            raise pa.ArrowInvalid(
-                f"Error converting to Arrow format: {str(e)}"
-            ) from e
-        except Exception as e:
-            logger.error(f"Unexpected error creating Arrow table: {str(e)}")
-            raise
+        Args:
+            dataset_id: Dataset ID to retrieve.
 
-        response = None
-        with ArizeFlightClient(
-            api_key=self._sdk_config.api_key,
-            host=self._sdk_config.flight_server_host,
-            port=self._sdk_config.flight_server_port,
-            scheme=self._sdk_config.flight_scheme,
-            request_verify=self._sdk_config.request_verify,
-            max_chunksize=self._sdk_config.pyarrow_max_chunksize,
-        ) as flight_client:
-            try:
-                response = flight_client.create_dataset(
-                    space_id=space_id,
-                    dataset_name=name,
-                    pa_table=pa_table,
-                )
-            except Exception as e:
-                msg = f"Error during update request: {str(e)}"
-                logger.error(msg)
-                raise RuntimeError(msg) from e
-        if response is None:
-            # This should not happen with proper Flight client implementation,
-            # but we handle it defensively
-            msg = "No response received from flight server during update"
-            logger.error(msg)
-            raise RuntimeError(msg)
-        # The response from flightserver is the dataset ID. To return the dataset
-        # object we make a GET query
-        dataset = self.get(dataset_id=response)
-        return dataset
+        Returns:
+            The dataset object.
 
-    def _list_examples(
+        Raises:
+            arize._generated.api_client.exceptions.ApiException: If the REST API
+                returns an error response (e.g. 401/403/404/429).
+        """
+        return self._api.datasets_get(dataset_id=dataset_id)
+
+    @prerelease_endpoint(key="datasets.delete", stage=ReleaseStage.BETA)
+    def delete(self, *, dataset_id: str) -> None:
+        """Delete a dataset by ID.
+
+        This operation is irreversible.
+
+        Args:
+            dataset_id: Dataset ID to delete.
+
+        Returns: This method returns None on success (common empty 204 response)
+
+        Raises:
+            arize._generated.api_client.exceptions.ApiException: If the REST API
+                returns an error response (e.g. 401/403/404/429).
+        """
+        return self._api.datasets_delete(dataset_id=dataset_id)
+
+    @prerelease_endpoint(key="datasets.list_examples", stage=ReleaseStage.BETA)
+    def list_examples(
         self,
         *,
         dataset_id: str,
         dataset_version_id: str = "",
         limit: int = 100,
         all: bool = False,
-    ):
+    ) -> models.DatasetsExamplesList200Response:
+        """List examples for a dataset (optionally for a specific version).
+
+        If `dataset_version_id` is not provided (empty string), the server selects
+        the latest dataset version.
+
+        Pagination notes:
+            - The response includes `pagination` for forward compatibility.
+            - Cursor pagination may not be fully implemented by the server yet.
+            - If `all=True`, this method retrieves all examples via the Flight path,
+              and returns them in a single response with `has_more=False`.
+
+        Args:
+            dataset_id: Dataset ID to list examples for.
+            dataset_version_id: Dataset version ID. If empty, the latest version is
+                selected.
+            limit: Maximum number of examples to return when `all=False`. The server
+                enforces an upper bound.
+            all: If True, fetch all examples (ignores `limit`) via Flight and return a
+                single response.
+
+        Returns:
+            A response object containing `examples` and `pagination` metadata.
+
+        Raises:
+            RuntimeError: If the Flight request fails or returns no response when
+                `all=True`.
+            arize._generated.api_client.exceptions.ApiException: If the REST API
+                returns an error response when `all=False` (e.g. 401/403/404/429).
+        """
         if not all:
-            return self._api.datasets_list_examples(
+            return self._api.datasets_examples_list(
                 dataset_id=dataset_id,
                 dataset_version_id=dataset_version_id,
                 limit=limit,
@@ -184,14 +274,17 @@ class DatasetsClient:
                 resource_updated_at=dataset_updated_at,
             )
         if dataset_df is not None:
-            return models.DatasetsListExamples200Response(
-                examples=dataset_df.to_dict(orient="records")
+            return models.DatasetsExamplesList200Response(
+                examples=dataset_df.to_dict(orient="records"),  # type: ignore
+                pagination=models.PaginationMetadata(
+                    has_more=False,  # Note that all=True
+                ),
             )
 
         with ArizeFlightClient(
             api_key=self._sdk_config.api_key,
-            host=self._sdk_config.flight_server_host,
-            port=self._sdk_config.flight_server_port,
+            host=self._sdk_config.flight_host,
+            port=self._sdk_config.flight_port,
             scheme=self._sdk_config.flight_scheme,
             request_verify=self._sdk_config.request_verify,
             max_chunksize=self._sdk_config.pyarrow_max_chunksize,
@@ -203,8 +296,8 @@ class DatasetsClient:
                     dataset_version_id=dataset_version_id,
                 )
             except Exception as e:
-                msg = f"Error during request: {str(e)}"
-                logger.error(msg)
+                msg = f"Error during request: {e!s}"
+                logger.exception(msg)
                 raise RuntimeError(msg) from e
         if dataset_df is None:
             # This should not happen with proper Flight client implementation,
@@ -222,12 +315,139 @@ class DatasetsClient:
             resource_data=dataset_df,
         )
 
-        return models.DatasetsListExamples200Response(
-            examples=dataset_df.to_dict(orient="records")
+        return models.DatasetsExamplesList200Response(
+            examples=dataset_df.to_dict(orient="records"),  # type: ignore
+            pagination=models.PaginationMetadata(
+                has_more=False,  # Note that all=True
+            ),
         )
 
+    # TODO(Kiko): Needs flightserver support
+    @prerelease_endpoint(
+        key="datasets.append_examples", stage=ReleaseStage.BETA
+    )
+    def append_examples(
+        self,
+        *,
+        dataset_id: str,
+        dataset_version_id: str = "",
+        examples: list[dict[str, object]] | pd.DataFrame,
+    ) -> models.Dataset:
+        """Append new examples to an existing dataset.
+
+        This method adds examples to an existing dataset version. If
+        `dataset_version_id` is not provided (empty string), the server appends
+        the examples to the latest dataset version.
+
+        The inserted examples are assigned system-generated IDs by the server.
+
+        Payload requirements (server-enforced):
+            - Each example may contain arbitrary user-defined fields.
+            - Do not include system-managed fields on input: `id`, `created_at`,
+              `updated_at` (requests containing these fields will be rejected).
+            - Each example must contain at least one property (i.e. empty
+              examples are not invalid).
+
+        Args:
+            dataset_id: Dataset ID to append examples to.
+            dataset_version_id: Optional dataset version ID to append examples to. If empty,
+                the latest dataset version is selected.
+            examples: Examples to append, provided as either:
+                - a list of JSON-like dicts, or
+                - a pandas DataFrame (converted to records before upload).
+
+        Returns:
+            The updated dataset object. To see the examples, use `list_examples()`.
+
+        Raises:
+            AssertionError: If `examples` is not a list of dicts or a pandas
+                DataFrame.
+            arize._generated.api_client.exceptions.ApiException: If the REST API
+                returns an error response (e.g. 400/401/403/404/429).
+        """
+        from arize._generated import api_client as gen
+
+        if not isinstance(examples, list | pd.DataFrame):
+            raise TypeError(
+                "Examples must be a list of dicts or a pandas DataFrame"
+            )
+
+        data = (
+            examples.to_dict(orient="records")  # type: ignore
+            if isinstance(examples, pd.DataFrame)
+            else examples
+        )
+        body = gen.DatasetsExamplesInsertRequest(examples=data)
+
+        return self._api.datasets_examples_insert(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_id,
+            datasets_examples_insert_request=body,
+        )
+
+    def _create_dataset_via_flight(
+        self,
+        name: str,
+        space_id: str,
+        examples: pd.DataFrame,
+    ) -> object:
+        """Internal method to create a dataset using Flight protocol for large example sets."""
+        data = examples.copy()
+        # Convert datetime columns to int64 (ms since epoch)
+        data = convert_datetime_columns_to_int(data)
+        data = convert_boolean_columns_to_str(data)
+        data = _set_default_columns_for_dataset(data)
+        data = convert_default_columns_to_json_str(data)
+
+        validation_errors = validate_dataset_df(data)
+        if validation_errors:
+            raise RuntimeError([e.error_message() for e in validation_errors])
+
+        # Convert to Arrow table
+        try:
+            logger.debug("Converting data to Arrow format")
+            pa_table = pa.Table.from_pandas(data, preserve_index=False)
+        except pa.ArrowInvalid as e:
+            logger.exception(INVALID_ARROW_CONVERSION_MSG)
+            raise pa.ArrowInvalid(
+                f"Error converting to Arrow format: {e!s}"
+            ) from e
+        except Exception:
+            logger.exception("Unexpected error creating Arrow table")
+            raise
+
+        response = None
+        with ArizeFlightClient(
+            api_key=self._sdk_config.api_key,
+            host=self._sdk_config.flight_host,
+            port=self._sdk_config.flight_port,
+            scheme=self._sdk_config.flight_scheme,
+            request_verify=self._sdk_config.request_verify,
+            max_chunksize=self._sdk_config.pyarrow_max_chunksize,
+        ) as flight_client:
+            try:
+                response = flight_client.create_dataset(
+                    space_id=space_id,
+                    dataset_name=name,
+                    pa_table=pa_table,
+                )
+            except Exception as e:
+                msg = f"Error during create request: {e!s}"
+                logger.exception(msg)
+                raise RuntimeError(msg) from e
+        if response is None:
+            # This should not happen with proper Flight client implementation,
+            # but we handle it defensively
+            msg = "No response received from flight server during update"
+            logger.error(msg)
+            raise RuntimeError(msg)
+        # The response from flightserver is the dataset ID. To return the dataset
+        # object we make a GET query
+        return self.get(dataset_id=response)
+
 
 def _set_default_columns_for_dataset(df: pd.DataFrame) -> pd.DataFrame:
+    """Set default values for created_at and updated_at columns if missing or null."""
     current_time = int(time.time() * 1000)
     if "created_at" in df.columns:
         if df["created_at"].isnull().values.any():  # type: ignore

arize/datasets/errors.py

@@ -1,21 +1,29 @@
+"""Dataset-specific exception classes."""
+
 from abc import ABC, abstractmethod
 
 
 class DatasetError(Exception, ABC):
+    """Base exception for dataset-related errors."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     @abstractmethod
     def __repr__(self) -> str:
-        pass
+        """Return a string representation for debugging and logging."""
 
     @abstractmethod
     def error_message(self) -> str:
-        pass
+        """Return the error message for this exception."""
 
 
 class InvalidSessionError(DatasetError):
+    """Raised when credentials are not provided or invalid."""
+
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             "Credentials not provided or invalid. Please pass in the correct api_key when "
             "initiating a new ArizeExportClient. Alternatively, you can set up credentials "
@@ -23,39 +31,61 @@ class InvalidSessionError(DatasetError):
         )
 
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "InvalidSessionError()"
 
 
 class InvalidConfigFileError(DatasetError):
+    """Raised when configuration file is invalid or misconfigured."""
+
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return "Invalid/Misconfigured Configuration File"
 
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "InvalidConfigFileError()"
 
 
 class IDColumnUniqueConstraintError(DatasetError):
+    """Raised when id column contains duplicate values."""
+
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return "'id' column must contain unique values"
 
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "IDColumnUniqueConstraintError()"
 
 
 class RequiredColumnsError(DatasetError):
+    """Raised when required columns are missing from the dataset."""
+
     def __init__(self, missing_columns: set) -> None:
+        """Initialize the exception with missing columns context.
+
+        Args:
+            missing_columns: Set of required columns that are missing.
+        """
         self.missing_columns = missing_columns
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return f"Missing required columns: {self.missing_columns}"
 
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return f"RequiredColumnsError({self.missing_columns})"
 
 
 class EmptyDatasetError(DatasetError):
+    """Raised when dataset DataFrame has no rows."""
+
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return "DataFrame must have at least one row in it."
 
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "EmptyDatasetError()"

arize/datasets/validation.py

@@ -1,4 +1,4 @@
-from typing import List
+"""Dataset validation logic for structure and content checks."""
 
 import pandas as pd
 
@@ -7,7 +7,17 @@ from arize.datasets import errors as err
 
 def validate_dataset_df(
     df: pd.DataFrame,
-) -> List[err.DatasetError]:
+) -> list[err.DatasetError]:
+    """Validate a dataset DataFrame for structural and content errors.
+
+    Checks for required columns, unique ID values, and non-empty data.
+
+    Args:
+        df: The pandas DataFrame to validate.
+
+    Returns:
+        A list of DatasetError objects found during validation. Empty list if valid.
+    """
     ## check all require columns are present
     required_columns_errors = _check_required_columns(df)
     if required_columns_errors:
@@ -19,14 +29,14 @@ def validate_dataset_df(
         return id_column_unique_constraint_error
 
     # check DataFrame has at least one row in it
-    emtpy_dataframe_error = _check_empty_dataframe(df)
-    if emtpy_dataframe_error:
-        return emtpy_dataframe_error
+    empty_dataframe_error = _check_empty_dataframe(df)
+    if empty_dataframe_error:
+        return empty_dataframe_error
 
     return []
 
 
-def _check_required_columns(df: pd.DataFrame) -> List[err.DatasetError]:
+def _check_required_columns(df: pd.DataFrame) -> list[err.DatasetError]:
     required_columns = ["id", "created_at", "updated_at"]
     missing_columns = set(required_columns) - set(df.columns)
     if missing_columns:
@@ -34,13 +44,13 @@ def _check_required_columns(df: pd.DataFrame) -> List[err.DatasetError]:
     return []
 
 
-def _check_id_column_is_unique(df: pd.DataFrame) -> List[err.DatasetError]:
+def _check_id_column_is_unique(df: pd.DataFrame) -> list[err.DatasetError]:
     if not df["id"].is_unique:
         return [err.IDColumnUniqueConstraintError()]
     return []
 
 
-def _check_empty_dataframe(df: pd.DataFrame) -> List[err.DatasetError]:
+def _check_empty_dataframe(df: pd.DataFrame) -> list[err.DatasetError]:
     if df.empty:
         return [err.EmptyDatasetError()]
     return []

arize/embeddings/__init__.py

@@ -1,3 +1,5 @@
+"""Embedding generation and use case utilities for the Arize SDK."""
+
 from arize.embeddings.auto_generator import EmbeddingGenerator
 from arize.embeddings.usecases import UseCases