arize 8.0.0a22__py3-none-any.whl → 8.0.0b0__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
@@ -11,56 +13,77 @@ from arize.config import SDKConfiguration
 if TYPE_CHECKING:
     from arize.datasets.client import DatasetsClient
     from arize.experiments.client import ExperimentsClient
-    from arize.models.client import MLModelsClient
+    from arize.ml.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): Clean commented lines over the SDK
+# TODO(Kiko): Implement https://github.com/Arize-ai/arize/pull/59917
+
+# TODO(Kiko): Go private connect. Need a `base_domain`, such that we get:
+# - api.<base_domain>
+# - app.<base_domain>
+# - flight.<base_domain>
+# - otlp.<base_domain>
+
+# 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, ML models, projects, and spans. It uses SDKConfiguration internally to
+    manage configuration settings.
+
+    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)
+
+    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
+            ... )
     """
 
-    _SUBCLIENTS = {
+    _SUBCLIENTS: ClassVar[dict[str, tuple[str, str]]] = {
         "datasets": (
             "arize.datasets.client",
             "DatasetsClient",
@@ -69,16 +92,20 @@ class ArizeClient(LazySubclientsMixin):
             "arize.experiments.client",
             "ExperimentsClient",
         ),
+        "projects": (
+            "arize.projects.client",
+            "ProjectsClient",
+        ),
+        "ml": (
+            "arize.ml.client",
+            "MLModelsClient",
+        ),
         "spans": (
             "arize.spans.client",
             "SpansClient",
         ),
-        "models": (
-            "arize.models.client",
-            "MLModelsClient",
-        ),
     }
-    _EXTRAS = {
+    _EXTRAS: ClassVar[dict[str, tuple[str | None, tuple[str, ...]]]] = {
         # Gate only the generated-backed ones
         "datasets": (
             "datasets-experiments",
@@ -114,31 +141,82 @@ class ArizeClient(LazySubclientsMixin):
         # Imports are gated in each method of the models client
         # This is to allow for very lean package install if people only
         # want to stream ML records
-        "models": (None, ()),
+        "ml": (None, ()),
     }
 
     def __init__(
         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 to this class 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 +225,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 +241,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")
+    def ml(self) -> MLModelsClient:
+        """Access the ML models client for ML model operations (lazy-loaded)."""
+        return self.__getattr__("ml")
 
     @property
-    def models(self) -> MLModelsClient:
-        return self.__getattr__("models")
+    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,
+        #     'ml': lazy,
+        #   }
+        # )
         lines = [f"{self.__class__.__name__}("]
         # Indent the SDKConfiguration repr
         cfg_repr = repr(self.sdk_config).splitlines()
@@ -205,7 +313,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(