arize 8.0.0b2__py3-none-any.whl → 8.0.1__py3-none-any.whl

arize/__init__.py

@@ -2,6 +2,7 @@
 
 import logging
 from collections.abc import Mapping
+from typing import Literal, cast
 
 from arize._generated.api_client import models
 from arize.client import ArizeClient
@@ -82,7 +83,11 @@ def make_to_df(field_name: str) -> object:
 
         # Drop None/NaN columns if requested
         if exclude_none in ("any", "all", True):
-            drop_how = "all" if exclude_none is True else exclude_none
+            drop_how: Literal["any", "all"] = (
+                "all"
+                if exclude_none is True
+                else cast("Literal['any', 'all']", exclude_none)
+            )
             df.dropna(axis=1, how=drop_how, inplace=True)
 
         if convert_dtypes:
@@ -92,6 +97,8 @@ def make_to_df(field_name: str) -> object:
     return to_df
 
 
+# Monkey-patch convenience methods onto generated response models
+# Type ignore comments needed: mypy can't verify runtime attribute additions
 models.DatasetsList200Response.to_df = make_to_df("datasets")  # type: ignore[attr-defined]
 models.DatasetsExamplesList200Response.to_df = make_to_df("examples")  # type: ignore[attr-defined]
 models.ExperimentsList200Response.to_df = make_to_df("experiments")  # type: ignore[attr-defined]

arize/_exporter/client.py

@@ -1,4 +1,3 @@
-# type: ignore[pb2]
 import logging
 from dataclasses import dataclass
 from datetime import datetime
@@ -41,7 +40,7 @@ class ArizeExportClient:
         batch_id: str = "",
         include_actuals: bool = False,
         stream_chunk_size: int | None = None,
-    ) -> object:
+    ) -> pd.DataFrame:
         """Exports data of a specific model in the Arize platform to a pandas dataframe.
 
         The export covers a defined time interval and model environment, and can
@@ -237,6 +236,22 @@ class ArizeExportClient:
         columns: list | None = None,
         stream_chunk_size: int | None = None,
     ) -> tuple[flight.FlightStreamReader | None, int]:
+        # Validate inputs first before creating logging context
+        validate_input_type(space_id, "space_id", str)
+        validate_input_type(model_id, "model_id", str)
+        validate_input_type(environment, "environment", Environments)
+        validate_input_type(include_actuals, "include_actuals", bool)
+        validate_input_type(start_time, "start_time", datetime)
+        validate_input_type(end_time, "end_time", datetime)
+        validate_input_type(model_version, "model_version", str)
+        validate_input_type(batch_id, "batch_id", str)
+        validate_input_type(where, "where", str)
+        validate_input_type(columns, "columns", list, allow_none=True)
+        validate_input_type(
+            stream_chunk_size, "stream_chunk_size", int, allow_none=True
+        )
+        validate_start_end_time(start_time, end_time)
+
         # Bind common context for this operation
         log = CtxAdapter(
             logger,
@@ -258,20 +273,6 @@ class ArizeExportClient:
             },
         )
         log.debug("Getting stream reader...")
-        validate_input_type(space_id, "space_id", str)
-        validate_input_type(model_id, "model_id", str)
-        validate_input_type(environment, "environment", Environments)
-        validate_input_type(include_actuals, "include_actuals", bool)
-        validate_input_type(start_time, "start_time", datetime)
-        validate_input_type(end_time, "end_time", datetime)
-        validate_input_type(model_version, "model_version", str)
-        validate_input_type(batch_id, "batch_id", str)
-        validate_input_type(where, "where", str)
-        validate_input_type(columns, "columns", list, allow_none=True)
-        validate_input_type(
-            stream_chunk_size, "stream_chunk_size", int, allow_none=True
-        )
-        validate_start_end_time(start_time, end_time)
 
         # Create query descriptor
         query_descriptor = flight_pb2.RecordQueryDescriptor(
@@ -300,7 +301,7 @@ class ArizeExportClient:
         try:
             flight_info = self.flight_client.get_flight_info(
                 flight.FlightDescriptor.for_command(
-                    json_format.MessageToJson(query_descriptor)  # type: ignore
+                    json_format.MessageToJson(query_descriptor)
                 ),
             )
             logger.info("Fetching data...")

arize/_exporter/parsers/tracing_data_parser.py

@@ -1,5 +1,7 @@
 import json
 import logging
+from collections.abc import Callable
+from typing import Any, TypeGuard
 
 import numpy as np
 import pandas as pd
@@ -28,7 +30,10 @@ logger = logging.getLogger(__name__)
 # of the error is on the data; It should not prevent a user from continuing to use the data
 class OtelTracingDataTransformer:
     def _apply_column_transformation(
-        self, df: pd.DataFrame, col_name: str, transform_func: object
+        self,
+        df: pd.DataFrame,
+        col_name: str,
+        transform_func: Callable[[Any], Any],
     ) -> str | None:
         """Apply a transformation to a column and return error message if it fails."""
         try:
@@ -89,7 +94,7 @@ class OtelTracingDataTransformer:
             if col.name in df.columns
         ]
         for col_name in dirty_string_column_names:
-            df[col_name] = df[col_name].apply(self._clean_json_string)
+            df[col_name] = df[col_name].apply(self._clean_json_string)  # type: ignore[arg-type]
 
         # Convert timestamp columns to datetime objects
         timestamp_column_names: list[str] = [
@@ -102,7 +107,7 @@ class OtelTracingDataTransformer:
         ]
         for col_name in timestamp_column_names:
             df[col_name] = df[col_name].apply(
-                self._convert_timestamp_to_datetime
+                self._convert_timestamp_to_datetime  # type: ignore[arg-type]
             )
 
         for err in errors:
@@ -138,7 +143,7 @@ class OtelTracingDataTransformer:
             return None
         return None
 
-    def _is_non_empty_string(self, value: object) -> bool:
+    def _is_non_empty_string(self, value: object) -> TypeGuard[str]:
         return isinstance(value, str) and value != ""
 
     def _deserialize_json_string_to_dict(self, value: str) -> object:

arize/_exporter/validation.py

@@ -35,7 +35,7 @@ def validate_input_value(
     if input in choices:
         return
     raise ValueError(
-        f"{input_name} is {input}, but must be one of {', '.join(choices)}"
+        f"{input_name} is {input}, but must be one of {', '.join(str(c) for c in choices)}"
     )
 
 

arize/_flight/client.py

@@ -1,11 +1,10 @@
-# type: ignore[pb2]
 from __future__ import annotations
 
 import base64
 import logging
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias
 
 from google.protobuf import json_format
 from pyarrow import flight
@@ -20,15 +19,15 @@ from arize.version import __version__
 
 if TYPE_CHECKING:
     import types
-    from collections.abc import Iterable
+    from collections.abc import Iterator
 
     import pandas as pd
     import pyarrow as pa
 
 
-BytesPair = tuple[bytes, bytes]
-Headers = list[BytesPair]
-FlightPostArrowFileResponse = (
+BytesPair: TypeAlias = tuple[bytes, bytes]
+Headers: TypeAlias = list[BytesPair]
+FlightPostArrowFileResponse: TypeAlias = (
     flight_pb2.WriteSpanEvaluationResponse
     | flight_pb2.WriteSpanAnnotationResponse
     | flight_pb2.WriteSpanAttributesMetadataResponse
@@ -90,7 +89,11 @@ class ArizeFlightClient:
     # ---------- Connection management ----------
 
     def _ensure_client(self) -> flight.FlightClient:
-        """Lazily initialize and return the underlying Flight client connection."""
+        """Lazily initialize and return the underlying Flight client connection.
+
+        Returns:
+            flight.FlightClient: The initialized Apache Arrow Flight client.
+        """
         client = object.__getattribute__(self, "_client")
         if client is not None:
             return client
@@ -135,7 +138,11 @@ class ArizeFlightClient:
     # ---------- methods simple passthrough wrappers ----------
 
     def get_flight_info(self, *args: object, **kwargs: object) -> object:
-        """Get flight information. Passthrough to underlying Flight client with auth options."""
+        """Get flight information. Passthrough to underlying Flight client with auth options.
+
+        Returns:
+            object: FlightInfo object containing metadata about the requested data stream.
+        """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.get_flight_info(*args, **kwargs)
@@ -146,6 +153,9 @@ class ArizeFlightClient:
         """Retrieve data stream via Flight DoGet.
 
         Passthrough to underlying Flight client with auth options.
+
+        Returns:
+            flight.FlightStreamReader: A stream reader for retrieving Arrow record batches.
         """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
@@ -153,10 +163,15 @@ class ArizeFlightClient:
 
     def do_put(
         self, *args: object, **kwargs: object
-    ) -> [flight.FlightStreamWriter, flight.FlightMetadataReader]:
+    ) -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
         """Upload data stream via Flight DoPut.
 
         Passthrough to underlying Flight client with auth options.
+
+        Returns:
+            tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]: A tuple containing
+                a stream writer for uploading Arrow record batches and a metadata reader for
+                receiving server responses.
         """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
@@ -164,10 +179,13 @@ class ArizeFlightClient:
 
     def do_action(
         self, *args: object, **kwargs: object
-    ) -> Iterable[flight.Result]:
+    ) -> Iterator[flight.Result]:
         """Execute an action via Flight DoAction.
 
         Passthrough to underlying Flight client with auth options.
+
+        Returns:
+            Iterable[flight.Result]: An iterable of Result objects from the action execution.
         """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
@@ -210,6 +228,10 @@ class ArizeFlightClient:
             FlightRequestType.ANNOTATION,
             FlightRequestType.METADATA,
         ):
+            if project_name is None:
+                raise ValueError(
+                    f"project_name is required for {request_type.name} request type"
+                )
             proto_schema = get_pb_schema_tracing(project_name=project_name)
             base64_schema = base64.b64encode(proto_schema.SerializeToString())
             pa_schema = append_to_pyarrow_metadata(
@@ -260,8 +282,6 @@ class ArizeFlightClient:
             case FlightRequestType.LOG_EXPERIMENT_DATA:
                 res = flight_pb2.PostExperimentDataResponse()
                 res.ParseFromString(flight_response.to_pybytes())
-            case _:
-                raise ValueError(f"Unsupported request_type: {request_type}")
         return res
 
     # ---------- dataset methods ----------
@@ -271,7 +291,7 @@ class ArizeFlightClient:
         space_id: str,
         dataset_name: str,
         pa_table: pa.Table,
-    ) -> str:
+    ) -> str | None:
         """Create a new dataset via Flight DoPut.
 
         Args:

arize/_lazy.py

@@ -52,7 +52,7 @@ class LazySubclientsMixin:
             # Determine which parameters this subclient needs
             # and build kwargs accordingly
             sig = inspect.signature(klass.__init__)
-            kwargs = {}
+            kwargs: dict[str, object] = {}
             if "sdk_config" in sig.parameters:
                 kwargs["sdk_config"] = self.sdk_config
             if "generated_client" in sig.parameters:
@@ -72,7 +72,14 @@ class OptionalDependencyError(ImportError): ...
 
 
 def _can_import(module_name: str) -> bool:
-    """Check if a module can be imported without raising an exception."""
+    """Check if a module can be imported without raising an exception.
+
+    Args:
+        module_name: The fully qualified module name to check (e.g., 'numpy', 'sklearn.preprocessing').
+
+    Returns:
+        bool: True if the module can be imported successfully, False otherwise.
+    """
     try:
         import_module(module_name)
     except Exception:
@@ -86,6 +93,18 @@ def require(
     required: tuple[str, ...],
     pkgname: str = "arize",
 ) -> None:
+    """Ensure required optional dependencies are installed, raising an error if missing.
+
+    Args:
+        extra_key: The extras group key for pip install (e.g., 'mimic', 'embeddings').
+            Used in the error message to guide users.
+        required: Tuple of required module names to check for availability.
+        pkgname: The package name for installation instructions. Defaults to 'arize'.
+
+    Raises:
+        OptionalDependencyError: If any of the required modules cannot be imported.
+            The error message includes pip install instructions with the extras group.
+    """
     if not required:
         return
     missing = [p for p in required if not _can_import(p)]
@@ -97,6 +116,22 @@ def require(
 
 
 def _dynamic_import(modname: str, retries: int = 2) -> types.ModuleType:
+    """Dynamically import a module with retry logic and sys.modules cleanup on failure.
+
+    Args:
+        modname: The fully qualified module name to import.
+        retries: Number of import attempts to make. Must be > 0. Defaults to 2.
+
+    Returns:
+        types.ModuleType: The successfully imported module.
+
+    Raises:
+        ValueError: If retries is <= 0.
+        ModuleNotFoundError: If the module cannot be found after all retry attempts.
+        ImportError: If the module import fails after all retry attempts.
+        KeyError: If a key error occurs during import after all retry attempts.
+    """
+
     def _attempt_import(remaining_attempts: int) -> types.ModuleType:
         try:
             return import_module(modname)

arize/client.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import logging
 import shutil
 from pathlib import Path
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, ClassVar, cast
 
 from arize._lazy import LazySubclientsMixin
 from arize.config import SDKConfiguration
@@ -20,20 +20,11 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-# 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 headers on each logging call
 # TODO(Kiko): InvalidAdditionalHeadersError is unused. Have we handled extra headers?
 
 # TODO(Kiko): Need to implement 'Update existing examples in a dataset'
 
-# TODO(Kiko): test caching in colab environment
 # TODO(Kiko): Protobuf versioning is too old
 # TODO(Kiko): Go through main APIs and add CtxAdapter where missing
 # TODO(Kiko): Search and handle other TODOs
@@ -128,6 +119,7 @@ class ArizeClient(LazySubclientsMixin):
         enable_caching: bool | None = None,
         single_host: str | None = None,
         single_port: int | None = None,
+        base_domain: str | None = None,
     ) -> None:
         """Initialize the Arize client with configuration parameters.
 
@@ -136,40 +128,72 @@ class ArizeClient(LazySubclientsMixin):
 
         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.
+                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.UNSET.
-            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.
+                overrides individual host/port settings.
+                ENV: ARIZE_REGION.
+                Default: Region.UNSET.
+            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.
+            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.
+            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.
+                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.
+                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.
+                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.
+                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.
+                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).
+            single_host: Single host for all endpoints. When specified, overrides
+                individual hosts.
+                ENV: ARIZE_SINGLE_HOST.
+                Default: None.
+            single_port: Single port for all endpoints. When specified, overrides
+                individual ports.
+                ENV: ARIZE_SINGLE_PORT.
+                Default: 0 (not set).
+            base_domain: Base domain for generating endpoint hosts as api.<base_domain>,
+                otlp.<base_domain>, flight.<base_domain>. Intended for Private Connect setups.
+                When specified, overrides individual hosts.
+                ENV: ARIZE_BASE_DOMAIN.
+                Default: None.
 
         Raises:
             MissingAPIKeyError: If api_key is not provided via argument or environment variable.
+            MultipleEndpointOverridesError: If multiple endpoint override options (region,
+                single_host/single_port, base_domain) are provided.
 
         Notes:
             Values provided to this class override environment variables, which in turn
@@ -213,6 +237,8 @@ class ArizeClient(LazySubclientsMixin):
             cfg_kwargs["single_host"] = single_host
         if single_port is not None:
             cfg_kwargs["single_port"] = single_port
+        if base_domain is not None:
+            cfg_kwargs["base_domain"] = base_domain
 
         # Only the explicitly provided fields are passed; the rest use
         # SDKConfiguration's default factories / defaults.
@@ -222,27 +248,27 @@ class ArizeClient(LazySubclientsMixin):
     @property
     def datasets(self) -> DatasetsClient:
         """Access the datasets client for dataset operations (lazy-loaded)."""
-        return self.__getattr__("datasets")
+        return cast("DatasetsClient", self.__getattr__("datasets"))
 
     @property
     def experiments(self) -> ExperimentsClient:
         """Access the experiments client for experiment operations (lazy-loaded)."""
-        return self.__getattr__("experiments")
+        return cast("ExperimentsClient", self.__getattr__("experiments"))
 
     @property
     def ml(self) -> MLModelsClient:
         """Access the ML models client for ML model operations (lazy-loaded)."""
-        return self.__getattr__("ml")
+        return cast("MLModelsClient", self.__getattr__("ml"))
 
     @property
     def projects(self) -> ProjectsClient:
         """Access the projects client for project operations (lazy-loaded)."""
-        return self.__getattr__("projects")
+        return cast("ProjectsClient", self.__getattr__("projects"))
 
     @property
     def spans(self) -> SpansClient:
         """Access the spans client for tracing and span operations (lazy-loaded)."""
-        return self.__getattr__("spans")
+        return cast("SpansClient", self.__getattr__("spans"))
 
     def __repr__(self) -> str:
         """Return a string representation of the Arize client configuration."""