arize 8.0.0b1__py3-none-any.whl → 8.0.0b4__py3-none-any.whl

arize/ml/stream_validation.py

@@ -1,7 +1,5 @@
 """Stream validation logic for ML model predictions."""
 
-# type: ignore[pb2]
-
 from arize.constants.ml import MAX_PREDICTION_ID_LEN, MIN_PREDICTION_ID_LEN
 from arize.exceptions.parameters import (
     InvalidValueType,
@@ -29,20 +27,8 @@ from arize.ml.types import (
 def validate_label(
     prediction_or_actual: str,
     model_type: ModelTypes,
-    label: str
-    | bool
-    | int
-    | float
-    | tuple[str | bool, float]
-    | ObjectDetectionLabel
-    | RankingPredictionLabel
-    | RankingActualLabel
-    | SemanticSegmentationLabel
-    | InstanceSegmentationPredictionLabel
-    | InstanceSegmentationActualLabel
-    | MultiClassPredictionLabel
-    | MultiClassActualLabel,
-    embedding_features: dict[str, Embedding],
+    label: object,
+    embedding_features: dict[str, Embedding] | None,
 ) -> None:
     """Validate a label value against the specified model type.
 
@@ -76,8 +62,17 @@ def validate_label(
 
 def _validate_numeric_label(
     model_type: ModelTypes,
-    label: str | bool | int | float | tuple[str | bool, float],
+    label: object,
 ) -> None:
+    """Validate that a label is numeric (int or float) for numeric model types.
+
+    Args:
+        model_type: The model type being validated.
+        label: The label value to validate.
+
+    Raises:
+        InvalidValueType: If the label is not an int or float.
+    """
     if not isinstance(label, (float, int)):
         raise InvalidValueType(
             f"label {label}",
@@ -88,8 +83,18 @@ def _validate_numeric_label(
 
 def _validate_categorical_label(
     model_type: ModelTypes,
-    label: str | bool | int | float | tuple[str | bool, float],
+    label: object,
 ) -> None:
+    """Validate that a label is categorical (scalar or tuple with confidence) for categorical model types.
+
+    Args:
+        model_type: The model type being validated.
+        label: The label value to validate.
+
+    Raises:
+        InvalidValueType: If the label is not a valid categorical type (bool, int, float, str,
+            or tuple of [str/bool, float]).
+    """
     is_valid = isinstance(label, (str, bool, int, float)) or (
         isinstance(label, tuple)
         and isinstance(label[0], (str, bool))
@@ -105,12 +110,20 @@ def _validate_categorical_label(
 
 def _validate_cv_label(
     prediction_or_actual: str,
-    label: ObjectDetectionLabel
-    | SemanticSegmentationLabel
-    | InstanceSegmentationPredictionLabel
-    | InstanceSegmentationActualLabel,
-    embedding_features: dict[str, Embedding],
+    label: object,
+    embedding_features: dict[str, Embedding] | None,
 ) -> None:
+    """Validate a computer vision label for object detection or segmentation tasks.
+
+    Args:
+        prediction_or_actual: Either 'prediction' or 'actual' to indicate label context.
+        label: The CV label to validate.
+        embedding_features: Dictionary of embedding features that must contain exactly one entry.
+
+    Raises:
+        InvalidValueType: If the label is not a valid CV label type.
+        ValueError: If embedding_features is None or doesn't contain exactly one feature.
+    """
     if (
         not isinstance(label, ObjectDetectionLabel)
         and not isinstance(label, SemanticSegmentationLabel)
@@ -138,8 +151,16 @@ def _validate_cv_label(
 
 
 def _validate_ranking_label(
-    label: RankingPredictionLabel | RankingActualLabel,
+    label: object,
 ) -> None:
+    """Validate a ranking label for ranking model types.
+
+    Args:
+        label: The ranking label to validate.
+
+    Raises:
+        InvalidValueType: If the label is not a RankingPredictionLabel or RankingActualLabel.
+    """
     if not isinstance(label, (RankingPredictionLabel, RankingActualLabel)):
         raise InvalidValueType(
             f"label {label}",
@@ -150,8 +171,16 @@ def _validate_ranking_label(
 
 
 def _validate_generative_llm_label(
-    label: str | bool | int | float,
+    label: object,
 ) -> None:
+    """Validate a label for generative LLM model types.
+
+    Args:
+        label: The label value to validate.
+
+    Raises:
+        InvalidValueType: If the label is not a bool, int, float, or str.
+    """
     is_valid = isinstance(label, (str, bool, int, float))
     if not is_valid:
         raise InvalidValueType(
@@ -162,8 +191,16 @@ def _validate_generative_llm_label(
 
 
 def _validate_multi_class_label(
-    label: MultiClassPredictionLabel | MultiClassActualLabel,
+    label: object,
 ) -> None:
+    """Validate a multi-class label for multi-class model types.
+
+    Args:
+        label: The multi-class label to validate.
+
+    Raises:
+        InvalidValueType: If the label is not a MultiClassPredictionLabel or MultiClassActualLabel.
+    """
     if not isinstance(
         label, (MultiClassPredictionLabel, MultiClassActualLabel)
     ):
@@ -185,7 +222,7 @@ def validate_and_convert_prediction_id(
     """Validate and convert a prediction ID to string format, or generate one if absent.
 
     Args:
-        prediction_id: The prediction ID to validate/convert, or None.
+        prediction_id: The prediction ID to validate/convert, or :obj:`None`.
         environment: The environment context (training, validation, production).
         prediction_label: Optional prediction label for delayed record detection.
         actual_label: Optional actual label for delayed record detection.

arize/ml/surrogate_explainer/mimic.py

@@ -19,6 +19,7 @@ from arize.ml.types import (
     CATEGORICAL_MODEL_TYPES,
     NUMERIC_MODEL_TYPES,
     ModelTypes,
+    _normalize_column_names,
 )
 
 if TYPE_CHECKING:
@@ -36,7 +37,7 @@ class Mimic:
         """Initialize the Mimic explainer with training data and model.
 
         Args:
-            X: Training data DataFrame for the surrogate model.
+            X: Training data :class:`pandas.DataFrame` for the surrogate model.
             model_func: Model function to explain.
         """
         self.explainer = MimicExplainer(
@@ -48,7 +49,7 @@ class Mimic:
         )
 
     def explain(self, X: pd.DataFrame) -> pd.DataFrame:
-        """Explain feature importance for the given input DataFrame."""
+        """Explain feature importance for the given input :class:`pandas.DataFrame`."""
         return pd.DataFrame(
             self.explainer.explain_local(X).local_importance_values,
             columns=X.columns,
@@ -59,8 +60,8 @@ class Mimic:
     def augment(
         df: pd.DataFrame, schema: Schema, model_type: ModelTypes
     ) -> tuple[pd.DataFrame, Schema]:
-        """Augment the DataFrame and schema with SHAP values for explainability."""
-        features = schema.feature_column_names
+        """Augment the :class:`pandas.DataFrame` and schema with SHAP values for explainability."""
+        features = _normalize_column_names(schema.feature_column_names)
         X = df[features]
 
         if X.shape[1] == 0:
@@ -85,25 +86,32 @@ class Mimic:
                 )
 
             # model func requires 1 positional argument
-            def model_func(_: object) -> object:  # type: ignore
+            def model_func(_: object) -> object:
                 return np.column_stack((1 - y, y))
 
         elif model_type in NUMERIC_MODEL_TYPES:
-            y_col_name = schema.prediction_label_column_name
+            y_col_name_nullable: str | None = (
+                schema.prediction_label_column_name
+            )
             if schema.prediction_score_column_name is not None:
-                y_col_name = schema.prediction_score_column_name
-            y = df[y_col_name].to_numpy()
+                y_col_name_nullable = schema.prediction_score_column_name
+            if y_col_name_nullable is None:
+                raise ValueError(
+                    f"For {model_type} models, either prediction_label_column_name "
+                    "or prediction_score_column_name must be specified"
+                )
+            y = df[y_col_name_nullable].to_numpy()
 
             _finite_count = np.isfinite(y).sum()
             if len(y) - _finite_count:
                 raise ValueError(
                     f"To calculate surrogate explainability for {model_type}, "
                     f"predictions must not contain NaN or infinite values, but "
-                    f"{len(y) - _finite_count} NaN or infinite value(s) are found in {y_col_name}."
+                    f"{len(y) - _finite_count} NaN or infinite value(s) are found in {y_col_name_nullable}."
                 )
 
             # model func requires 1 positional argument
-            def model_func(_: object) -> object:  # type: ignore
+            def model_func(_: object) -> object:
                 return y
 
         else:

arize/ml/types.py

@@ -2,47 +2,47 @@
 
 import logging
 import math
+import sys
 from collections.abc import Iterator
 from dataclasses import asdict, dataclass, replace
 from datetime import datetime
 from decimal import Decimal
 from enum import Enum, unique
 from itertools import chain
-from typing import (
-    NamedTuple,
-    Self,
-)
+from typing import NamedTuple
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
 
 import numpy as np
 
 from arize.constants.ml import (
-    # MAX_MULTI_CLASS_NAME_LENGTH,
-    # MAX_NUMBER_OF_MULTI_CLASS_CLASSES,
     MAX_MULTI_CLASS_NAME_LENGTH,
     MAX_NUMBER_OF_MULTI_CLASS_CLASSES,
     MAX_NUMBER_OF_SIMILARITY_REFERENCES,
     MAX_RAW_DATA_CHARACTERS,
     MAX_RAW_DATA_CHARACTERS_TRUNCATION,
-    # MAX_RAW_DATA_CHARACTERS,
-    # MAX_RAW_DATA_CHARACTERS_TRUNCATION,
 )
 from arize.exceptions.parameters import InvalidValueType
-
-#
-# from arize.utils.constants import (
-#     MAX_MULTI_CLASS_NAME_LENGTH,
-#     MAX_NUMBER_OF_MULTI_CLASS_CLASSES,
-#     MAX_NUMBER_OF_SIMILARITY_REFERENCES,
-#     MAX_RAW_DATA_CHARACTERS,
-#     MAX_RAW_DATA_CHARACTERS_TRUNCATION,
-# )
-# from arize.utils.errors import InvalidValueType
 from arize.logging import get_truncation_warning_message
 from arize.utils.types import is_dict_of, is_iterable_of, is_list_of
 
 logger = logging.getLogger(__name__)
 
 
+def _normalize_column_names(
+    col_names: "list[str] | TypedColumns | None",
+) -> list[str]:
+    """Convert TypedColumns or list to a flat list of column names."""
+    if col_names is None:
+        return []
+    if isinstance(col_names, list):
+        return col_names
+    return col_names.get_all_column_names()
+
+
 @unique
 class ModelTypes(Enum):
     """Enum representing supported model types in Arize."""
@@ -204,7 +204,7 @@ class Embedding(NamedTuple):
             )
         # Fail if not all elements in list are floats
         allowed_types = (int, float, np.int16, np.int32, np.float16, np.float32)
-        if not all(isinstance(val, allowed_types) for val in self.vector):  # type: ignore
+        if not all(isinstance(val, allowed_types) for val in self.vector):
             raise TypeError(
                 f"Embedding vector must be a vector of integers and/or floats. Got "
                 f"{emb_name}.vector = {self.vector}"
@@ -283,7 +283,7 @@ class Embedding(NamedTuple):
 
     @staticmethod
     def _is_valid_iterable(
-        data: str | list[str] | list[float] | np.ndarray,
+        data: object,
     ) -> bool:
         """Validates that the input data field is of the correct iterable type.
 
@@ -299,30 +299,6 @@ class Embedding(NamedTuple):
         return any(isinstance(data, t) for t in (list, np.ndarray))
 
 
-# @dataclass
-# class _PromptOrResponseText:
-#     data: str
-#
-#     def validate(self, name: str) -> None:
-#         # Validate that data is a string
-#         if not isinstance(self.data, str):
-#             raise TypeError(f"'{name}' must be a str")
-#
-#         character_count = len(self.data)
-#         if character_count > MAX_RAW_DATA_CHARACTERS:
-#             raise ValueError(
-#                 f"'{name}' field must not contain more than {MAX_RAW_DATA_CHARACTERS} characters. "
-#                 f"Found {character_count}."
-#             )
-#         elif character_count > MAX_RAW_DATA_CHARACTERS_TRUNCATION:
-#             logger.warning(
-#                 get_truncation_warning_message(
-#                     f"'{name}'", MAX_RAW_DATA_CHARACTERS_TRUNCATION
-#                 )
-#             )
-#         return None
-
-
 class LLMRunMetadata(NamedTuple):
     """Metadata for LLM execution including token counts and latency."""
 
@@ -1021,22 +997,6 @@ class LLMRunMetadataColumnNames:
         )
 
 
-# @dataclass
-# class DocumentColumnNames:
-#     id_column_name: Optional[str] = None
-#     version_column_name: Optional[str] = None
-#     text_embedding_column_names: Optional[EmbeddingColumnNames] = None
-#
-#     def __iter__(self):
-#         return iter(
-#             (
-#                 self.id_column_name,
-#                 self.version_column_name,
-#                 self.text_embedding_column_names,
-#             )
-#         )
-#
-#
 @dataclass
 class SimilarityReference:
     """Reference to a prediction for similarity search operations."""
@@ -1250,7 +1210,7 @@ class Schema(BaseSchema):
     actual_score_column_name: str | None = None
     shap_values_column_names: dict[str, str] | None = None
     embedding_feature_column_names: dict[str, EmbeddingColumnNames] | None = (
-        None  # type:ignore
+        None
     )
     prediction_group_id_column_name: str | None = None
     rank_column_name: str | None = None
@@ -1268,7 +1228,7 @@ class Schema(BaseSchema):
     prompt_template_column_names: PromptTemplateColumnNames | None = None
     llm_config_column_names: LLMConfigColumnNames | None = None
     llm_run_metadata_column_names: LLMRunMetadataColumnNames | None = None
-    retrieved_document_ids_column_name: list[str] | None = None
+    retrieved_document_ids_column_name: str | None = None
     multi_class_threshold_scores_column_name: str | None = None
     semantic_segmentation_prediction_column_names: (
         SemanticSegmentationColumnNames | None
@@ -1285,7 +1245,7 @@ class Schema(BaseSchema):
 
     def get_used_columns_counts(self) -> dict[str, int]:
         """Return a dict mapping column names to their usage count."""
-        columns_used_counts = {}
+        columns_used_counts: dict[str, int] = {}
 
         for field in self.__dataclass_fields__:
             if field.endswith("column_name"):
@@ -1294,7 +1254,7 @@ class Schema(BaseSchema):
                     add_to_column_count_dictionary(columns_used_counts, col)
 
         if self.feature_column_names is not None:
-            for col in self.feature_column_names:
+            for col in _normalize_column_names(self.feature_column_names):
                 add_to_column_count_dictionary(columns_used_counts, col)
 
         if self.embedding_feature_column_names is not None:
@@ -1313,7 +1273,7 @@ class Schema(BaseSchema):
                     )
 
         if self.tag_column_names is not None:
-            for col in self.tag_column_names:
+            for col in _normalize_column_names(self.tag_column_names):
                 add_to_column_count_dictionary(columns_used_counts, col)
 
         if self.shap_values_column_names is not None:
@@ -1458,7 +1418,7 @@ class CorpusSchema(BaseSchema):
 
     def get_used_columns_counts(self) -> dict[str, int]:
         """Return a dict mapping column names to their usage count."""
-        columns_used_counts = {}
+        columns_used_counts: dict[str, int] = {}
 
         if self.document_id_column_name is not None:
             add_to_column_count_dictionary(
@@ -1531,7 +1491,7 @@ def add_to_column_count_dictionary(
 
     Args:
         column_dictionary: Dictionary mapping column names to counts.
-        col: The column name to increment, or None to skip.
+        col: The column name to increment, or :obj:`None` to skip.
     """
     if col:
         if col in column_dictionary:

arize/pre_releases.py

@@ -3,14 +3,15 @@
 import functools
 import logging
 from collections.abc import Callable
-from enum import StrEnum
+from enum import Enum
+from typing import TypeVar, cast
 
 from arize.version import __version__
 
 logger = logging.getLogger(__name__)
 
 
-class ReleaseStage(StrEnum):
+class ReleaseStage(Enum):
     """Enum representing the release stage of API features."""
 
     ALPHA = "alpha"
@@ -19,19 +20,21 @@ class ReleaseStage(StrEnum):
 
 _WARNED: set[str] = set()
 
+_F = TypeVar("_F", bound=Callable)
+
 
 def _format_prerelease_message(*, key: str, stage: ReleaseStage) -> str:
     article = "an" if stage is ReleaseStage.ALPHA else "a"
     return (
-        f"[{stage.upper()}] {key} is {article} {stage} API "
+        f"[{stage.value.upper()}] {key} is {article} {stage.value} API "
         f"in Arize SDK v{__version__} and may change without notice."
     )
 
 
-def prerelease_endpoint(*, stage: ReleaseStage, key: str) -> object:
+def prerelease_endpoint(*, key: str, stage: ReleaseStage) -> Callable[[_F], _F]:
     """Decorate a method to emit a prerelease warning via logging once per process."""
 
-    def deco(fn: Callable[..., object]) -> object:
+    def deco(fn: _F) -> _F:
         @functools.wraps(fn)
         def wrapper(*args: object, **kwargs: object) -> object:
             if key not in _WARNED:
@@ -39,6 +42,7 @@ def prerelease_endpoint(*, stage: ReleaseStage, key: str) -> object:
                 logger.warning(_format_prerelease_message(key=key, stage=stage))
             return fn(*args, **kwargs)
 
-        return wrapper
+        # Cast: functools.wraps preserves function signature at runtime but mypy can't verify this
+        return cast("_F", wrapper)
 
     return deco

arize/projects/client.py

@@ -9,6 +9,7 @@ from arize.pre_releases import ReleaseStage, prerelease_endpoint
 
 if TYPE_CHECKING:
     from arize._generated.api_client import models
+    from arize._generated.api_client.api_client import ApiClient
     from arize.config import SDKConfiguration
 
 logger = logging.getLogger(__name__)
@@ -26,18 +27,21 @@ class ProjectsClient:
     :class:`arize.config.SDKConfiguration`.
     """
 
-    def __init__(self, *, sdk_config: SDKConfiguration) -> None:
+    def __init__(
+        self, *, sdk_config: SDKConfiguration, generated_client: ApiClient
+    ) -> None:
         """
         Args:
             sdk_config: Resolved SDK configuration.
+            generated_client: Shared generated API client instance.
         """  # noqa: D205, D212
         self._sdk_config = sdk_config
 
         # 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.ProjectsApi(self._sdk_config.get_generated_client())
+        # Use the provided client directly
+        self._api = gen.ProjectsApi(generated_client)
 
     @prerelease_endpoint(key="projects.list", stage=ReleaseStage.BETA)
     def list(
@@ -125,7 +129,8 @@ class ProjectsClient:
         Args:
             project_id: Project ID.
 
-        Returns: This method returns None on success (common empty 204 response)
+        Returns:
+            This method returns None on success (common empty 204 response).
 
         Raises:
             arize._generated.api_client.exceptions.ApiException: If the API request fails

arize/regions.py

@@ -1,19 +1,19 @@
 """Region definitions and configuration for Arize deployment zones."""
 
 from dataclasses import dataclass
-from enum import StrEnum
+from enum import Enum
 
 from arize.constants.config import DEFAULT_FLIGHT_PORT
 
 
-class Region(StrEnum):
+class Region(Enum):
     """Enum representing available Arize deployment regions."""
 
-    US_CENTRAL_1 = "us-central-1a"
-    EU_WEST_1 = "eu-west-1a"
-    CA_CENTRAL_1 = "ca-central-1a"
-    US_EAST_1 = "us-east-1b"
-    UNSPECIFIED = ""
+    CA_CENTRAL_1A = "ca-central-1a"
+    EU_WEST_1A = "eu-west-1a"
+    US_CENTRAL_1A = "us-central-1a"
+    US_EAST_1B = "us-east-1b"
+    UNSET = ""
 
 
 @dataclass(frozen=True)
@@ -28,13 +28,13 @@ class RegionEndpoints:
 
 def _get_region_endpoints(region: Region) -> RegionEndpoints:
     return RegionEndpoints(
-        api_host=f"api.{region}.arize.com",
-        otlp_host=f"otlp.{region}.arize.com",
-        flight_host=f"flight.{region}.arize.com",
+        api_host=f"api.{region.value}.arize.com",
+        otlp_host=f"otlp.{region.value}.arize.com",
+        flight_host=f"flight.{region.value}.arize.com",
         flight_port=DEFAULT_FLIGHT_PORT,
     )
 
 
 REGION_ENDPOINTS: dict[Region, RegionEndpoints] = {
-    r: _get_region_endpoints(r) for r in Region if r != Region.UNSPECIFIED
+    r: _get_region_endpoints(r) for r in list(Region) if r != Region.UNSET
 }