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

arize/_lazy.py

@@ -5,7 +5,7 @@ import logging
 import sys
 import threading
 from importlib import import_module
-from typing import TYPE_CHECKING, Any, Dict, Tuple
+from typing import TYPE_CHECKING, ClassVar
 
 if TYPE_CHECKING:
     from arize.config import SDKConfiguration
@@ -14,15 +14,15 @@ logger = logging.getLogger(__name__)
 
 
 class LazySubclientsMixin:
-    _SUBCLIENTS: Dict[str, Tuple[str, str]] = {}
-    _EXTRAS: Dict[str, Tuple[str | None, Tuple[str, ...]]] = {}
+    _SUBCLIENTS: ClassVar[dict[str, tuple[str, str]]] = {}
+    _EXTRAS: ClassVar[dict[str, tuple[str | None, tuple[str, ...]]]] = {}
 
-    def __init__(self, sdk_config: SDKConfiguration):
+    def __init__(self, sdk_config: SDKConfiguration) -> None:
         self.sdk_config = sdk_config
-        self._lazy_cache: Dict[str, Any] = {}
+        self._lazy_cache: dict[str, object] = {}
         self._lazy_lock = threading.Lock()
 
-    def __getattr__(self, name: str) -> Any:
+    def __getattr__(self, name: str) -> object:
         subs = self._SUBCLIENTS
         if name not in subs:
             raise AttributeError(
@@ -50,26 +50,31 @@ class LazySubclientsMixin:
             self._lazy_cache[name] = instance
             return instance
 
-    def __dir__(self):
+    def __dir__(self) -> list[str]:
         return sorted({*super().__dir__(), *self._SUBCLIENTS.keys()})
 
 
 class OptionalDependencyError(ImportError): ...
 
 
+def _can_import(module_name: str) -> bool:
+    """Check if a module can be imported without raising an exception."""
+    try:
+        import_module(module_name)
+    except Exception:
+        return False
+    else:
+        return True
+
+
 def require(
     extra_key: str | None,
-    required: Tuple[str, ...],
-    pkgname="arize",
-):
+    required: tuple[str, ...],
+    pkgname: str = "arize",
+) -> None:
     if not required:
         return
-    missing = []
-    for p in required:
-        try:
-            import_module(p)
-        except Exception:
-            missing.append(p)
+    missing = [p for p in required if not _can_import(p)]
     if missing:
         raise OptionalDependencyError(
             f"Missing optional dependencies: {', '.join(missing)}. "
@@ -77,11 +82,14 @@ def require(
         )
 
 
-def _dynamic_import(modname: str, retries: int = 2):
-    for i in range(retries):
+def _dynamic_import(modname: str, retries: int = 2) -> object:
+    def _attempt_import(remaining_attempts: int) -> object:
         try:
             return import_module(modname)
         except (ModuleNotFoundError, ImportError, KeyError):
             sys.modules.pop(modname, None)
-            if i + 1 == retries:
+            if remaining_attempts <= 1:
                 raise
+            return _attempt_import(remaining_attempts - 1)
+
+    return _attempt_import(retries) if retries > 0 else None

arize/client.py

@@ -1,9 +1,11 @@
+"""Main Arize SDK client for interacting with Arize AI platform services."""
+
 from __future__ import annotations
 
 import logging
 import shutil
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
 
 from arize._lazy import LazySubclientsMixin
 from arize.config import SDKConfiguration
@@ -12,55 +14,166 @@ if TYPE_CHECKING:
     from arize.datasets.client import DatasetsClient
     from arize.experiments.client import ExperimentsClient
     from arize.models.client import MLModelsClient
+    from arize.projects.client import ProjectsClient
+    from arize.regions import Region
     from arize.spans.client import SpansClient
 
 logger = logging.getLogger(__name__)
 
 
+# TODO(Kiko): Enforce type checking, remove all type ignores
+
+# TODO(Kiko): Go over docstrings
+# TODO(Kiko): Missing parameter descriptions in some docstrings
+# TODO(Kiko): Missing return descriptions in some docstrings
+
+# TODO(Kiko): Go over headers on each logging call
+
+# TODO(Kiko): Need to implement 'Update existing examples in a dataset'
+
 # TODO(Kiko): why logs don't show on scripts, only on jupyter notebooks
 # TODO(Kiko): test caching in colab environment
-# TODO(Kiko): models need to follow resource first pattern
-# -  models.DatasetsList200Response
-# -  models.DatasetsListExamples200Response
-# -  models.ExperimentsList200Response
-# -  models.ExperimentsRunsList200Response
 # TODO(Kiko): Protobuf versioning is too old
-# TODO(Kiko): It does not make any sense to require space ID in run_experiment, dataset ID should suffice
-# TODO(Kiko): Should probably wrap every single method of gen clients so that we can add nice docstrings
-# TODO(Kiko): Add flight max_chunksize opt to write_table. In config?
-# TODO(Kiko): experimental/datasets must be adapted into the datasets subclient
-# TODO(Kiko): experimental/prompt hub is missing
-# TODO(Kiko): exporter/utils/schema_parser is missing
 # TODO(Kiko): Go through main APIs and add CtxAdapter where missing
 # TODO(Kiko): Search and handle other TODOs
 # TODO(Kiko): Go over **every file** and do not import anything at runtime, use `if TYPE_CHECKING`
 # with `from __future__ import annotations` (must include for Python < 3.11)
-# TODO(Kiko): Go over docstrings
-# TODO(Kiko): Go headers on each logging call
-# TODO(Kiko): Make sure the client has same options as SDKConfiguration
+
+
 class ArizeClient(LazySubclientsMixin):
-    """
-    Root client for the Arize SDK. All parameters are optional. If not provided, they will be read
-    from environment variables. If those are absent, built-in defaults will be used with the only
-    exception of `api_key`, which is required to be provided either via argument or environment
-    variable (ARIZE_API_KEY).
-
-    Parameters (all optional):
-        api_key: If not provided, read from environment (ARIZE_API_KEY). If absent, raises.
-        api_host: If not provided, read from environment (ARIZE_API_HOST) or default.
-        api_scheme: If not provided, read from environment (ARIZE_API_INSECURE -> http/https) or default.
-        flight_server_host: If not provided, read from environment (ARIZE_FLIGHT_HOST) or default.
-        flight_server_port: If not provided, read from environment (ARIZE_FLIGHT_PORT) or default.
-        flight_scheme: If not provided, read from environment (ARIZE_FLIGHT_TRANSPORT_SCHEME) or default.
-        request_verify: If not provided, read from environment (ARIZE_REQUEST_VERIFY) or defaults to True.
-
-    Resolution order for each field:
-        1) Value passed here (if not None)
-        2) Environment variable (handled by SDKConfiguration class)
-        3) Built-in default constant (handled by SDKConfiguration class)
+    """Root client for the Arize SDK.
+
+    The ArizeClient provides access to all Arize platform services including datasets,
+    experiments, models, projects, and spans. It uses SDKConfiguration internally to
+    manage configuration settings.
+
+    Configuration Precedence
+    ------------------------
+    All parameters are optional (except api_key which must be provided via argument
+    or environment variable). For each parameter, values are resolved in this order:
+        1. Explicit value passed to constructor (highest priority)
+        2. Environment variable value (see SDKConfiguration for variable names)
+        3. Built-in default value (lowest priority)
+
+    Parameters
+    ----------
+    api_key : str
+        Arize API key for authentication. Required - must be provided via this parameter
+        or the ARIZE_API_KEY environment variable.
+        Environment variable: ARIZE_API_KEY
+        Default: None (raises MissingAPIKeyError if not provided)
+
+    region : Region, optional
+        Arize region (e.g., Region.US_CENTRAL, Region.EU_WEST). When specified,
+        overrides individual host/port settings with region-specific endpoints.
+        Environment variable: ARIZE_REGION
+        Default: Region.UNSPECIFIED
+
+    api_host : str, optional
+        Custom API endpoint host override.
+        Environment variable: ARIZE_API_HOST
+        Default: "api.arize.com"
+
+    api_scheme : str, optional
+        API endpoint scheme (http/https).
+        Environment variable: ARIZE_API_SCHEME
+        Default: "https"
+
+    otlp_host : str, optional
+        OTLP (OpenTelemetry Protocol) endpoint host override.
+        Environment variable: ARIZE_OTLP_HOST
+        Default: "otlp.arize.com"
+
+    otlp_scheme : str, optional
+        OTLP endpoint scheme (http/https).
+        Environment variable: ARIZE_OTLP_SCHEME
+        Default: "https"
+
+    flight_host : str, optional
+        Apache Arrow Flight endpoint host override.
+        Environment variable: ARIZE_FLIGHT_HOST
+        Default: "flight.arize.com"
+
+    flight_port : int, optional
+        Apache Arrow Flight endpoint port (1-65535).
+        Environment variable: ARIZE_FLIGHT_PORT
+        Default: 443
+
+    flight_scheme : str, optional
+        Apache Arrow Flight endpoint scheme.
+        Environment variable: ARIZE_FLIGHT_SCHEME
+        Default: "grpc+tls"
+
+    pyarrow_max_chunksize : int, optional
+        Maximum chunk size for PyArrow operations (1 to MAX_CHUNKSIZE).
+        Environment variable: ARIZE_MAX_CHUNKSIZE
+        Default: 10_000
+
+    request_verify : bool, optional
+        Whether to verify SSL certificates for HTTP requests.
+        Environment variable: ARIZE_REQUEST_VERIFY
+        Default: True
+
+    stream_max_workers : int, optional
+        Maximum number of worker threads for streaming operations (minimum: 1).
+        Environment variable: ARIZE_STREAM_MAX_WORKERS
+        Default: 8
+
+    stream_max_queue_bound : int, optional
+        Maximum queue size for streaming operations (minimum: 1).
+        Environment variable: ARIZE_STREAM_MAX_QUEUE_BOUND
+        Default: 5000
+
+    max_http_payload_size_mb : float, optional
+        Maximum HTTP payload size in megabytes (minimum: 1).
+        Environment variable: ARIZE_MAX_HTTP_PAYLOAD_SIZE_MB
+        Default: 100
+
+    arize_directory : str, optional
+        Directory for Arize SDK files (cache, logs, etc.).
+        Environment variable: ARIZE_DIRECTORY
+        Default: "~/.arize"
+
+    enable_caching : bool, optional
+        Whether to enable local caching.
+        Environment variable: ARIZE_ENABLE_CACHING
+        Default: True
+
+    single_host : str, optional
+        Single host to use for all endpoints. Overrides individual host settings.
+        Environment variable: ARIZE_SINGLE_HOST
+        Default: None (not set)
+
+    single_port : int, optional
+        Single port to use for all endpoints. Overrides individual port settings (0-65535).
+        Environment variable: ARIZE_SINGLE_PORT
+        Default: 0 (not set)
+
+    Examples:
+    --------
+    Initialize with API key only (other settings use defaults):
+
+        >>> client = ArizeClient(api_key="your-api-key")
+
+    Initialize with custom endpoints:
+
+        >>> client = ArizeClient(
+        ...     api_key="your-api-key",
+        ...     api_host="custom.api.com",
+        ...     flight_port=8443,
+        ... )
+
+    Initialize with region (overrides host/port settings):
+
+        >>> from arize.regions import Region
+        >>> client = ArizeClient(api_key="your-api-key", region=Region.EU_WEST)
+
+    See Also:
+    --------
+    SDKConfiguration : Internal configuration class with detailed parameter documentation
     """
 
-    _SUBCLIENTS = {
+    _SUBCLIENTS: ClassVar[dict[str, tuple[str, str]]] = {
         "datasets": (
             "arize.datasets.client",
             "DatasetsClient",
@@ -69,16 +182,20 @@ class ArizeClient(LazySubclientsMixin):
             "arize.experiments.client",
             "ExperimentsClient",
         ),
-        "spans": (
-            "arize.spans.client",
-            "SpansClient",
+        "projects": (
+            "arize.projects.client",
+            "ProjectsClient",
         ),
         "models": (
             "arize.models.client",
             "MLModelsClient",
         ),
+        "spans": (
+            "arize.spans.client",
+            "SpansClient",
+        ),
     }
-    _EXTRAS = {
+    _EXTRAS: ClassVar[dict[str, tuple[str | None, tuple[str, ...]]]] = {
         # Gate only the generated-backed ones
         "datasets": (
             "datasets-experiments",
@@ -121,24 +238,75 @@ class ArizeClient(LazySubclientsMixin):
         self,
         *,
         api_key: str | None = None,
+        region: Region | None = None,
         api_host: str | None = None,
         api_scheme: str | None = None,
         otlp_host: str | None = None,
         otlp_scheme: str | None = None,
-        flight_server_host: str | None = None,
-        flight_server_port: int | None = None,
+        flight_host: str | None = None,
+        flight_port: int | None = None,
         flight_scheme: str | None = None,
         pyarrow_max_chunksize: int | None = None,
         request_verify: bool | None = None,
         stream_max_workers: int | None = None,
         stream_max_queue_bound: int | None = None,
         max_http_payload_size_mb: float | None = None,
-        arize_direcory: str | None = None,
+        arize_directory: str | None = None,
         enable_caching: bool | None = None,
-    ):
+        single_host: str | None = None,
+        single_port: int | None = None,
+    ) -> None:
+        """Initialize the Arize client with configuration parameters.
+
+        All parameters are optional and follow the precedence order documented in the
+        class docstring: explicit argument > environment variable > default value.
+
+        Args:
+            api_key: Arize API key for authentication. Required - must be provided here
+                or via ARIZE_API_KEY environment variable. Raises MissingAPIKeyError if not set.
+            region: Arize region (e.g., Region.US_CENTRAL, Region.EU_WEST). When specified,
+                overrides individual host/port settings. ENV: ARIZE_REGION. Default: Region.UNSPECIFIED.
+            api_host: Custom API endpoint host. ENV: ARIZE_API_HOST. Default: "api.arize.com".
+            api_scheme: API endpoint scheme (http/https). ENV: ARIZE_API_SCHEME. Default: "https".
+            otlp_host: OTLP endpoint host. ENV: ARIZE_OTLP_HOST. Default: "otlp.arize.com".
+            otlp_scheme: OTLP endpoint scheme (http/https). ENV: ARIZE_OTLP_SCHEME. Default: "https".
+            flight_host: Apache Arrow Flight endpoint host. ENV: ARIZE_FLIGHT_HOST.
+                Default: "flight.arize.com".
+            flight_port: Apache Arrow Flight endpoint port (1-65535). ENV: ARIZE_FLIGHT_PORT.
+                Default: 443.
+            flight_scheme: Apache Arrow Flight endpoint scheme. ENV: ARIZE_FLIGHT_SCHEME.
+                Default: "grpc+tls".
+            pyarrow_max_chunksize: Maximum PyArrow chunk size (1 to MAX_CHUNKSIZE).
+                ENV: ARIZE_MAX_CHUNKSIZE. Default: 10_000.
+            request_verify: Whether to verify SSL certificates. ENV: ARIZE_REQUEST_VERIFY.
+                Default: True.
+            stream_max_workers: Maximum worker threads for streaming (minimum: 1).
+                ENV: ARIZE_STREAM_MAX_WORKERS. Default: 8.
+            stream_max_queue_bound: Maximum queue size for streaming (minimum: 1).
+                ENV: ARIZE_STREAM_MAX_QUEUE_BOUND. Default: 5000.
+            max_http_payload_size_mb: Maximum HTTP payload size in MB (minimum: 1).
+                ENV: ARIZE_MAX_HTTP_PAYLOAD_SIZE_MB. Default: 100.
+            arize_directory: Directory for SDK files (cache, logs, etc.).
+                ENV: ARIZE_DIRECTORY. Default: "~/.arize".
+            enable_caching: Whether to enable local caching. ENV: ARIZE_ENABLE_CACHING.
+                Default: True.
+            single_host: Single host for all endpoints (overrides individual hosts).
+                ENV: ARIZE_SINGLE_HOST. Default: None.
+            single_port: Single port for all endpoints (overrides individual ports).
+                ENV: ARIZE_SINGLE_PORT. Default: 0 (not set).
+
+        Raises:
+            MissingAPIKeyError: If api_key is not provided via argument or environment variable.
+
+        Notes:
+            Values provided here override environment variables, which in turn override
+            default values. See SDKConfiguration for detailed parameter documentation.
+        """
         cfg_kwargs: dict = {}
         if api_key is not None:
             cfg_kwargs["api_key"] = api_key
+        if region is not None:
+            cfg_kwargs["region"] = region
         if api_host is not None:
             cfg_kwargs["api_host"] = api_host
         if api_scheme is not None:
@@ -147,10 +315,10 @@ class ArizeClient(LazySubclientsMixin):
             cfg_kwargs["otlp_host"] = otlp_host
         if otlp_scheme is not None:
             cfg_kwargs["otlp_scheme"] = otlp_scheme
-        if flight_server_host is not None:
-            cfg_kwargs["flight_server_host"] = flight_server_host
-        if flight_server_port is not None:
-            cfg_kwargs["flight_server_port"] = flight_server_port
+        if flight_host is not None:
+            cfg_kwargs["flight_host"] = flight_host
+        if flight_port is not None:
+            cfg_kwargs["flight_port"] = flight_port
         if flight_scheme is not None:
             cfg_kwargs["flight_scheme"] = flight_scheme
         if pyarrow_max_chunksize is not None:
@@ -163,33 +331,63 @@ class ArizeClient(LazySubclientsMixin):
             cfg_kwargs["stream_max_queue_bound"] = stream_max_queue_bound
         if max_http_payload_size_mb is not None:
             cfg_kwargs["max_http_payload_size_mb"] = max_http_payload_size_mb
-        if arize_direcory is not None:
-            cfg_kwargs["arize_direcory"] = arize_direcory
+        if arize_directory is not None:
+            cfg_kwargs["arize_directory"] = arize_directory
         if enable_caching is not None:
             cfg_kwargs["enable_caching"] = enable_caching
+        if single_host is not None:
+            cfg_kwargs["single_host"] = single_host
+        if single_port is not None:
+            cfg_kwargs["single_port"] = single_port
 
         # Only the explicitly provided fields are passed; the rest use
-        # SDKConfiguration’s default factories / defaults.
+        # SDKConfiguration's default factories / defaults.
         super().__init__(SDKConfiguration(**cfg_kwargs))
 
     # typed properties for IDE completion
     @property
     def datasets(self) -> DatasetsClient:
+        """Access the datasets client for dataset operations (lazy-loaded)."""
         return self.__getattr__("datasets")
 
     @property
     def experiments(self) -> ExperimentsClient:
+        """Access the experiments client for experiment operations (lazy-loaded)."""
         return self.__getattr__("experiments")
 
-    @property
-    def spans(self) -> SpansClient:
-        return self.__getattr__("spans")
-
     @property
     def models(self) -> MLModelsClient:
+        """Access the models client for ML model operations (lazy-loaded)."""
         return self.__getattr__("models")
 
+    @property
+    def projects(self) -> ProjectsClient:
+        """Access the projects client for project operations (lazy-loaded)."""
+        return self.__getattr__("projects")
+
+    @property
+    def spans(self) -> SpansClient:
+        """Access the spans client for tracing and span operations (lazy-loaded)."""
+        return self.__getattr__("spans")
+
     def __repr__(self) -> str:
+        """Return a string representation of the Arize client configuration."""
+        # The repr looks like:
+        #  ArizeClient(
+        #   sdk_config=SDKConfiguration(
+        #     api_key='cacaca***',
+        #     api_host='api.arize.com',
+        #     ...
+        #     arize_directory='~/.arize',
+        #     enable_caching=True,
+        #   )
+        #   subclients={
+        #     'datasets': lazy,
+        #     'experiments': lazy,
+        #     'spans': lazy,
+        #     'models': lazy,
+        #   }
+        # )
         lines = [f"{self.__class__.__name__}("]
         # Indent the SDKConfiguration repr
         cfg_repr = repr(self.sdk_config).splitlines()
@@ -205,7 +403,22 @@ class ArizeClient(LazySubclientsMixin):
         return "\n".join(lines)
 
     def clear_cache(self) -> None:
-        p = Path(self.sdk_config.cache_dir)
+        """Clear the local cache directory.
+
+        Removes all cached data from the SDK's cache directory. This can be useful
+        when troubleshooting caching issues or freeing up disk space. The cache
+        directory is automatically recreated on subsequent operations that require
+        caching.
+
+        Raises:
+            NotADirectoryError: If the cache path exists but is not a directory.
+
+        Notes:
+            - This operation permanently deletes all cached data
+            - If the cache directory doesn't exist, a warning is logged but no error is raised
+            - The cache directory location is configured via sdk_config.cache_dir
+        """
+        p = Path(self.sdk_config.cache_dir).expanduser().resolve()
 
         if not p.exists():
             logger.warning(