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

arize/experiments/types.py

@@ -1,18 +1,16 @@
+"""Type definitions and data models for experiments."""
+
 from __future__ import annotations
 
 import json
 import textwrap
+from collections.abc import Awaitable, Callable, Iterable, Mapping
 from copy import copy, deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from importlib.metadata import version
 from random import getrandbits
 from typing import (
-    Any,
-    Awaitable,
-    Callable,
-    Iterable,
-    Mapping,
     cast,
 )
 
@@ -35,8 +33,8 @@ TraceId = str
 
 @dataclass(frozen=True)
 class Example:
-    """
-    Represents an example in an experiment dataset.
+    """Represents an example in an experiment dataset.
+
     Args:
         id: The unique identifier for the example.
         updated_at: The timestamp when the example was last updated.
@@ -54,6 +52,7 @@ class Example:
     dataset_row: Mapping[str, JSONSerializable] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
+        """Initialize example fields from dataset_row if provided."""
         if self.dataset_row is not None:
             object.__setattr__(
                 self, "dataset_row", _make_read_only(self.dataset_row)
@@ -90,7 +89,8 @@ class Example:
             object.__setattr__(self, "metadata", self.metadata)
 
     @classmethod
-    def from_dict(cls, obj: Mapping[str, Any]) -> Example:
+    def from_dict(cls, obj: Mapping[str, object]) -> Example:
+        """Create an Example instance from a dictionary."""
         return cls(
             id=obj["id"],
             input=obj["input"],
@@ -100,6 +100,7 @@ class Example:
         )
 
     def __repr__(self) -> str:
+        """Return a formatted string representation of the example."""
         spaces = " " * 4
         name = self.__class__.__name__
         identifiers = [f'{spaces}id="{self.id}",']
@@ -123,7 +124,9 @@ class Example:
         return "\n".join([f"{name}(", *identifiers, *contents, ")"])
 
 
-def _shorten(obj: Any, width: int = 50) -> Any:
+def _shorten(
+    obj: dict[str, object] | list[object] | str | object, width: int = 50
+) -> dict[str, object] | list[object] | str | object:
     if isinstance(obj, str):
         return textwrap.shorten(obj, width=width, placeholder="...")
     if isinstance(obj, dict):
@@ -135,7 +138,9 @@ def _shorten(obj: Any, width: int = 50) -> Any:
     return obj
 
 
-def _make_read_only(obj: Any) -> Any:
+def _make_read_only(
+    obj: dict[str, object] | list[object] | str | object,
+) -> dict[str, object] | list[object] | str | object:
     if isinstance(obj, dict):
         return _ReadOnly({k: _make_read_only(v) for k, v in obj.items()})
     if isinstance(obj, str):
@@ -146,25 +151,25 @@ def _make_read_only(obj: Any) -> Any:
 
 
 class _ReadOnly(ObjectProxy):  # type: ignore[misc]
-    def __setitem__(self, *args: Any, **kwargs: Any) -> Any:
+    def __setitem__(self, *args: object, **kwargs: object) -> object:
         raise NotImplementedError
 
-    def __delitem__(self, *args: Any, **kwargs: Any) -> Any:
+    def __delitem__(self, *args: object, **kwargs: object) -> object:
         raise NotImplementedError
 
-    def __iadd__(self, *args: Any, **kwargs: Any) -> Any:
+    def __iadd__(self, *args: object, **kwargs: object) -> object:
         raise NotImplementedError
 
-    def pop(self, *args: Any, **kwargs: Any) -> Any:
+    def pop(self, *args: object, **kwargs: object) -> object:
         raise NotImplementedError
 
-    def append(self, *args: Any, **kwargs: Any) -> Any:
+    def append(self, *args: object, **kwargs: object) -> object:
         raise NotImplementedError
 
-    def __copy__(self, *args: Any, **kwargs: Any) -> Any:
+    def __copy__(self, *args: object, **kwargs: object) -> object:
         return copy(self.__wrapped__)
 
-    def __deepcopy__(self, *args: Any, **kwargs: Any) -> Any:
+    def __deepcopy__(self, *args: object, **kwargs: object) -> object:
         return deepcopy(self.__wrapped__)
 
     def __repr__(self) -> str:
@@ -180,6 +185,8 @@ def _blue(text: str) -> str:
 
 @dataclass(frozen=True)
 class TestCase:
+    """Container for an experiment test case with example data and repetition number."""
+
     example: Example
     repetition_number: RepetitionNumber
 
@@ -194,8 +201,8 @@ def _exp_id() -> str:
 
 @dataclass(frozen=True)
 class ExperimentRun:
-    """
-    Represents a single run of an experiment.
+    """Represents a single run of an experiment.
+
     Args:
         start_time: The start time of the experiment run.
         end_time: The end time of the experiment run.
@@ -219,7 +226,8 @@ class ExperimentRun:
     trace_id: TraceId | None = None
 
     @classmethod
-    def from_dict(cls, obj: Mapping[str, Any]) -> ExperimentRun:
+    def from_dict(cls, obj: Mapping[str, object]) -> ExperimentRun:
+        """Create an ExperimentRun instance from a dictionary."""
         return cls(
             start_time=obj["start_time"],
             end_time=obj["end_time"],
@@ -233,6 +241,11 @@ class ExperimentRun:
         )
 
     def __post_init__(self) -> None:
+        """Validate that exactly one of output or error is specified.
+
+        Raises:
+            ValueError: If both or neither output and error are specified.
+        """
         if (self.output is None) == (self.error is None):
             raise ValueError(
                 "Must specify exactly one of experiment_run_output or error"
@@ -241,8 +254,8 @@ class ExperimentRun:
 
 @dataclass(frozen=True)
 class ExperimentEvaluationRun:
-    """
-    Represents a single evaluation run of an experiment.
+    """Represents a single evaluation run of an experiment.
+
     Args:
         experiment_run_id: The unique identifier for the experiment run.
         start_time: The start time of the evaluation run.
@@ -266,7 +279,8 @@ class ExperimentEvaluationRun:
     trace_id: TraceId | None = None
 
     @classmethod
-    def from_dict(cls, obj: Mapping[str, Any]) -> ExperimentEvaluationRun:
+    def from_dict(cls, obj: Mapping[str, object]) -> ExperimentEvaluationRun:
+        """Create an ExperimentEvaluationRun instance from a dictionary."""
         return cls(
             experiment_run_id=obj["experiment_run_id"],
             start_time=obj["start_time"],
@@ -280,6 +294,11 @@ class ExperimentEvaluationRun:
         )
 
     def __post_init__(self) -> None:
+        """Validate that exactly one of result or error is specified.
+
+        Raises:
+            ValueError: If both or neither result and error are specified.
+        """
         if bool(self.result) == bool(self.error):
             raise ValueError("Must specify either result or error")
 
@@ -288,6 +307,11 @@ _LOCAL_TIMEZONE = datetime.now(timezone.utc).astimezone().tzinfo
 
 
 def local_now() -> datetime:
+    """Get the current datetime in the local timezone.
+
+    Returns:
+        A datetime object representing the current time in local timezone.
+    """
     return datetime.now(timezone.utc).astimezone(tz=_LOCAL_TIMEZONE)
 
 
@@ -303,10 +327,12 @@ class _HasStats:
 
     def __str__(self) -> str:
         try:
-            assert int(version("pandas").split(".")[0]) >= 1
+            pandas_major = int(version("pandas").split(".")[0])
+            if pandas_major < 1:
+                raise ImportError("Pandas version < 1.0")  # noqa: TRY301
             # `tabulate` is used by pandas >= 1.0 in DataFrame.to_markdown()
             import tabulate  # noqa: F401
-        except (AssertionError, ImportError):
+        except ImportError:
             text = self.stats.__str__()
         else:
             text = self.stats.to_markdown(index=False)
@@ -315,8 +341,7 @@ class _HasStats:
 
 @dataclass(frozen=True)
 class _TaskSummary(_HasStats):
-    """
-    Summary statistics of experiment task executions.
+    """Summary statistics of experiment task executions.
 
     **Users should not instantiate this object directly.**
     """
@@ -344,7 +369,7 @@ class _TaskSummary(_HasStats):
             "n_runs": n_runs,
             "n_errors": n_errors,
             **(
-                dict(top_error=_top_string(df.loc[:, "error"]))
+                {"top_error": _top_string(df.loc[:, "error"])}
                 if n_errors
                 else {}
             ),
@@ -355,12 +380,12 @@ class _TaskSummary(_HasStats):
         return summary
 
     @classmethod
-    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
+    def __new__(cls, *args: object, **kwargs: object) -> object:
         # Direct instantiation by users is discouraged.
         raise NotImplementedError
 
     @classmethod
-    def __init_subclass__(cls, **kwargs: Any) -> None:
+    def __init_subclass__(cls, **kwargs: object) -> None:
         # Direct sub-classing by users is discouraged.
         raise NotImplementedError
 
@@ -368,7 +393,7 @@ class _TaskSummary(_HasStats):
 def _top_string(s: pd.Series, length: int = 100) -> str | None:
     if (cnt := s.dropna().str.slice(0, length).value_counts()).empty:
         return None
-    return cast(str, cnt.sort_values(ascending=False).index[0])
+    return cast("str", cnt.sort_values(ascending=False).index[0])
 
 
 @dataclass

arize/logging.py

@@ -1,10 +1,13 @@
+"""Logging utilities and configuration for the Arize SDK."""
+
 from __future__ import annotations
 
 import json
 import logging
 import os
 import sys
-from typing import Any, Dict, Iterable, Mapping
+from collections.abc import Iterable, Mapping
+from typing import Any, ClassVar
 
 from arize.config import _parse_bool
 from arize.constants.config import (
@@ -30,7 +33,10 @@ _LEVEL_MAP = {
 class CtxAdapter(logging.LoggerAdapter):
     """LoggerAdapter that merges bound context with per-call extras safely."""
 
-    def process(self, msg, kwargs):
+    def process(
+        self, msg: object, kwargs: dict[str, object]
+    ) -> tuple[object, dict[str, object]]:
+        """Process the logging call by merging bound and call extras."""
         call_extra = _coerce_mapping(kwargs.pop("extra", None))
         bound_extra = _coerce_mapping(self.extra)
         merged = (
@@ -42,7 +48,7 @@ class CtxAdapter(logging.LoggerAdapter):
             kwargs["extra"] = merged
         return msg, kwargs
 
-    def with_extra(self, **more) -> CtxAdapter:
+    def with_extra(self, **more: object) -> CtxAdapter:
         """Return a copy of this adapter with additional bound extras."""
         base = _coerce_mapping(self.extra)
         base.update(_coerce_mapping(more))
@@ -54,12 +60,7 @@ class CtxAdapter(logging.LoggerAdapter):
 
 
 class CustomLogFormatter(logging.Formatter):
-    grey = "\x1b[38;21m"
-    blue = "\x1b[38;5;39m"
-    yellow = "\x1b[33m"
-    red = "\x1b[38;5;196m"
-    bold_red = "\x1b[31;1m"
-    reset = "\x1b[0m"
+    """Custom log formatter with color-coded output based on log level."""
 
     GREY = "\x1b[38;21m"
     BLUE = "\x1b[38;5;39m"
@@ -68,7 +69,7 @@ class CustomLogFormatter(logging.Formatter):
     BOLD_RED = "\x1b[31;1m"
     RESET = "\x1b[0m"
 
-    COLORS = {
+    COLORS: ClassVar[dict[int, str]] = {
         logging.DEBUG: BLUE,
         logging.INFO: GREY,
         logging.WARNING: YELLOW,
@@ -76,10 +77,16 @@ class CustomLogFormatter(logging.Formatter):
         logging.CRITICAL: BOLD_RED,
     }
 
-    def __init__(self, fmt: str):
+    def __init__(self, fmt: str) -> None:
+        """Initialize the colored formatter with a format string.
+
+        Args:
+            fmt: Format string for log messages.
+        """
         super().__init__(fmt=fmt)
 
     def format(self, record: logging.LogRecord) -> str:
+        """Format the log record with color based on log level."""
         # Build the base message without any color.
         base = super().format(record)
 
@@ -106,7 +113,7 @@ class JsonFormatter(logging.Formatter):
     """Minimal JSON formatter (one JSON object per line)."""
 
     # fields to skip copying from record.__dict__
-    _skip = {
+    _skip: ClassVar[set[str]] = {
         # "name",
         # "msg",
         # "args",
@@ -130,7 +137,8 @@ class JsonFormatter(logging.Formatter):
     }
 
     def format(self, record: logging.LogRecord) -> str:
-        payload: Dict[str, Any] = {
+        """Format the log record as a JSON string."""
+        payload: dict[str, object] = {
             # "time": self.formatTime(record, datefmt="%Y-%m-%dT%H:%M:%S%z"),
             # "logger": record.name,
             # "level": record.levelname,
@@ -156,12 +164,11 @@ def _parse_level(val: str | None, default: int = logging.INFO) -> int:
 
 
 def auto_configure_from_env() -> None:
-    """
-    If ARIZE_LOG is truthy, configure logging for 'arize' once,
-    using ARIZE_LOG_LEVEL and ARIZE_LOG_STRUCTURED if provided.
+    """If ARIZE_LOG is truthy, configure logging for 'arize' once.
+
+    Using ARIZE_LOG_LEVEL and ARIZE_LOG_STRUCTURED if provided.
     Otherwise, do nothing (library stays quiet with NullHandler).
     """
-
     if not _parse_bool(os.getenv(ENV_LOG_ENABLE, DEFAULT_LOG_ENABLE)):
         return
 
@@ -201,7 +208,16 @@ _STANDARD_RECORD_KEYS = {
 }
 
 
-def get_truncation_warning_message(instance, limit) -> str:
+def get_truncation_warning_message(instance: str, limit: int) -> str:
+    """Generate a warning message for data that exceeds character limits.
+
+    Args:
+        instance: The data instance type that may be truncated.
+        limit: The character limit threshold.
+
+    Returns:
+        A formatted warning message string.
+    """
     return (
         f"Attention: {instance} exceeding the {limit} character limit will be "
         "automatically truncated upon ingestion into the Arize platform. Should you require "
@@ -213,8 +229,7 @@ def configure_logging(
     level: int = logging.INFO,
     structured: bool = False,
 ) -> None:
-    """
-    Configure logging for the 'arize' logger.
+    """Configure logging for the 'arize' logger.
 
     Args:
         level: logging level (e.g., logging.INFO, logging.DEBUG)
@@ -241,6 +256,15 @@ def configure_logging(
 
 
 def log_a_list(values: Iterable[Any] | None, join_word: str) -> str:
+    """Format a list of values into a human-readable string with a joining word.
+
+    Args:
+        values: An iterable of values to format, or None.
+        join_word: The word to use before the last item (e.g., "and", "or").
+
+    Returns:
+        A formatted string joining the values, or empty string if no values.
+    """
     if values is None:
         return ""
     list_of_str = list(values)
@@ -253,13 +277,21 @@ def log_a_list(values: Iterable[Any] | None, join_word: str) -> str:
     )
 
 
-def get_arize_project_url(response: Any):
+def get_arize_project_url(response: object) -> str:
+    """Extract the Arize project URL from an API response.
+
+    Args:
+        response: The HTTP response object from an Arize API call.
+
+    Returns:
+        The real-time ingestion URI if present, otherwise an empty string.
+    """
     if "realTimeIngestionUri" in json.loads(response.content.decode()):
         return json.loads(response.content.decode())["realTimeIngestionUri"]
     return ""
 
 
-def _coerce_mapping(obj: Any) -> Dict[str, Any]:
+def _coerce_mapping(obj: object) -> dict[str, object]:
     """Return a shallow dict copy if obj is a Mapping[str, Any], else {}."""
     if isinstance(obj, Mapping):
         # force keys to str to satisfy logging's expectation

arize/models/__init__.py

@@ -0,0 +1 @@
+"""Model logging and validation functionality for the Arize SDK."""

arize/models/batch_validation/__init__.py

@@ -0,0 +1 @@
+"""Batch validation utilities for model data."""