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

arize/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility functions and helper modules for the Arize SDK."""

arize/utils/arrow.py

@@ -1,3 +1,5 @@
+"""Apache Arrow utilities for data serialization and file operations."""
+
 # type: ignore[pb2]
 from __future__ import annotations
 
@@ -5,7 +7,7 @@ import base64
 import logging
 import os
 import tempfile
-from typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa
 
@@ -23,16 +25,30 @@ def post_arrow_table(
     files_url: str,
     pa_table: pa.Table,
     proto_schema: pb2.Schema,
-    headers: Dict[str, str],
+    headers: dict[str, str],
     timeout: float | None,
     verify: bool,
     max_chunksize: int,
     tmp_dir: str = "",
 ) -> requests.Response:
-    # We import here to avoid depending onn requests for all arrow utils
+    """Post a PyArrow table to Arize via HTTP file upload.
+
+    Args:
+        files_url: The URL endpoint for file uploads.
+        pa_table: The PyArrow table containing the data.
+        proto_schema: The protobuf schema for the data.
+        headers: HTTP headers for the request.
+        timeout: Request timeout in seconds, or None for no timeout.
+        verify: Whether to verify SSL certificates.
+        max_chunksize: Maximum chunk size for splitting large tables.
+        tmp_dir: Temporary directory for serialization. Defaults to "".
+
+    Returns:
+        The HTTP response from the upload request.
+    """
+    # We import here to avoid depending on requests for all arrow utils
     import requests
 
-    logger.debug("Preparing to log Arrow table via file upload")
     logger.debug(
         "Preparing to log Arrow table via file upload",
         extra={"rows": pa_table.num_rows, "cols": pa_table.num_columns},
@@ -94,20 +110,20 @@ def post_arrow_table(
                 tdir.cleanup()  # cleaning the entire dir, no need to clean the file
             except Exception as e:
                 logger.warning(
-                    f"Failed to remove temporary directory {tdir.name}: {str(e)}"
+                    f"Failed to remove temporary directory {tdir.name}: {e!s}"
                 )
         elif cleanup_file:
             try:
                 os.remove(outfile)
             except Exception as e:
                 logger.warning(
-                    f"Failed to remove temporary file {outfile}: {str(e)}"
+                    f"Failed to remove temporary file {outfile}: {e!s}"
                 )
 
 
 def _append_to_pyarrow_metadata(
-    pa_schema: pa.Schema, new_metadata: Dict[str, Any]
-):
+    pa_schema: pa.Schema, new_metadata: dict[str, Any]
+) -> object:
     # Ensure metadata is handled correctly, even if initially None.
     metadata = pa_schema.metadata
     if metadata is None:
@@ -129,9 +145,10 @@ def _append_to_pyarrow_metadata(
 def _write_arrow_file(
     path: str, pa_table: pa.Table, pa_schema: pa.Schema, max_chunksize: int
 ) -> None:
-    with pa.OSFile(path, mode="wb") as sink, pa.ipc.RecordBatchStreamWriter(
-        sink, pa_schema
-    ) as writer:
+    with (
+        pa.OSFile(path, mode="wb") as sink,
+        pa.ipc.RecordBatchStreamWriter(sink, pa_schema) as writer,
+    ):
         writer.write_table(pa_table, max_chunksize)
 
 
@@ -145,10 +162,9 @@ def _maybe_log_project_url(response: requests.Response) -> None:
 
 
 def _mktemp_in(directory: str) -> str:
-    """
-    Create a unique temp file path inside `directory` without leaving
-    an open file descriptor around (Windows-safe). The file exists on
-    disk and is closed; caller can open/write it later.
+    """Create a unique temp file path inside `directory` without leaving an open file descriptor.
+
+    Windows-safe. The file exists on disk and is closed; caller can open/write it later.
     """
     with tempfile.NamedTemporaryFile(
         dir=directory,

arize/utils/cache.py

@@ -1,10 +1,16 @@
+"""Caching utilities for resource management and persistence."""
+
 from __future__ import annotations
 
 import logging
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pandas as pd
 
+if TYPE_CHECKING:
+    from datetime import datetime
+
 logger = logging.getLogger(__name__)
 
 
@@ -12,9 +18,21 @@ def load_cached_resource(
     cache_dir: str,
     resource: str,
     resource_id: str,
-    resource_updated_at: str | None,
+    resource_updated_at: datetime | None,
     format: str = "parquet",
 ) -> pd.DataFrame | None:
+    """Load a cached resource from the local cache directory.
+
+    Args:
+        cache_dir: Directory path for cache storage.
+        resource: Resource type name (e.g., "dataset", "experiment").
+        resource_id: Unique identifier for the resource.
+        resource_updated_at: Optional timestamp of last resource update.
+        format: File format for cached data. Defaults to "parquet".
+
+    Returns:
+        The cached DataFrame if found and valid, None otherwise.
+    """
     key = _get_cache_key(resource, resource_id, resource_updated_at)
     filepath = _get_abs_file_path(cache_dir, f"{key}.{format}", resource)
     if not filepath.exists():
@@ -30,10 +48,20 @@ def cache_resource(
     cache_dir: str,
     resource: str,
     resource_id: str,
-    resource_updated_at: str | None,
+    resource_updated_at: datetime | None,
     resource_data: pd.DataFrame,
     format: str = "parquet",
 ) -> None:
+    """Save a resource to the local cache directory.
+
+    Args:
+        cache_dir: Directory path for cache storage.
+        resource: Resource type name (e.g., "dataset", "experiment").
+        resource_id: Unique identifier for the resource.
+        resource_updated_at: Optional timestamp of last resource update.
+        resource_data: DataFrame containing the resource data.
+        format: File format for cached data. Defaults to "parquet".
+    """
     key = _get_cache_key(resource, resource_id, resource_updated_at)
     filepath = _get_abs_file_path(cache_dir, f"{key}.{format}", resource)
     filepath.parent.mkdir(parents=True, exist_ok=True)
@@ -44,12 +72,12 @@ def cache_resource(
 def _get_cache_key(
     resource: str,
     resource_id: str,
-    resource_updated_at: str | None,
+    resource_updated_at: datetime | None,
 ) -> str:
     # include updated_at if present to produce a new key when dataset changes
     key = f"{resource}_{resource_id}"
     if resource_updated_at:
-        key += f"_{resource_updated_at}"
+        key += f"_{resource_updated_at.strftime('%Y%m%dT%H%M%S')}"
     return key
 
 
@@ -58,8 +86,8 @@ def _get_abs_file_path(
     filename: str,
     subdirectory: str | None = None,
 ) -> Path:
-    """
-    Return an absolute path to a file located under `directory[/subdirectory]/filename`.
+    """Return an absolute path to a file located under `directory[/subdirectory]/filename`.
+
     Expands '~' and resolves relative components.
     """
     base = Path(directory).expanduser()

arize/utils/dataframe.py

@@ -1,5 +1,6 @@
+"""DataFrame manipulation and validation utilities."""
+
 import re
-from typing import List
 
 import pandas as pd
 
@@ -8,6 +9,11 @@ from arize.types import BaseSchema
 
 # Resets the dataframe index if it is not a RangeIndex
 def reset_dataframe_index(dataframe: pd.DataFrame) -> None:
+    """Reset the DataFrame index in-place if it is not a RangeIndex.
+
+    Args:
+        dataframe: The pandas DataFrame to reset.
+    """
     if not isinstance(dataframe.index, pd.RangeIndex):
         drop = dataframe.index.name in dataframe.columns
         dataframe.reset_index(inplace=True, drop=drop)
@@ -16,9 +22,20 @@ def reset_dataframe_index(dataframe: pd.DataFrame) -> None:
 def remove_extraneous_columns(
     df: pd.DataFrame,
     schema: BaseSchema | None = None,
-    column_list: List[str] | None = None,
+    column_list: list[str] | None = None,
     regex: str | None = None,
 ) -> pd.DataFrame:
+    """Filter DataFrame to keep only relevant columns based on schema, list, or regex.
+
+    Args:
+        df: The pandas DataFrame to filter.
+        schema: Optional schema defining used columns. Defaults to None.
+        column_list: Optional explicit list of columns to keep. Defaults to None.
+        regex: Optional regex pattern to match column names. Defaults to None.
+
+    Returns:
+        A filtered DataFrame containing only the relevant columns.
+    """
     relevant_columns = set()
     if schema is not None:
         relevant_columns.update(schema.get_used_columns())

arize/utils/online_tasks/__init__.py

@@ -1,3 +1,5 @@
+"""Online task processing utilities for dataframe preprocessing."""
+
 from arize.utils.online_tasks.dataframe_preprocessor import (
     extract_nested_data_to_column,
 )

arize/utils/online_tasks/dataframe_preprocessor.py

@@ -1,6 +1,7 @@
+"""DataFrame preprocessing utilities for online tasks."""
+
 import json
 import logging
-from typing import Any, List, Tuple
 
 import numpy as np
 import pandas as pd
@@ -8,17 +9,34 @@ import pandas as pd
 logger = logging.getLogger(__name__)
 
 
+class ColumnNotFoundError(Exception):
+    """Raised when a specified column is not found in the DataFrame."""
+
+    def __init__(self, attribute: str) -> None:
+        """Initialize with the attribute that couldn't be mapped to a column.
+
+        Args:
+            attribute: The attribute string that has no matching column prefix.
+        """
+        self.attribute = attribute
+        super().__init__(
+            f"No column found in DataFrame for attribute: {attribute}"
+        )
+
+
 def extract_nested_data_to_column(
-    attributes: List[str], df: pd.DataFrame
+    attributes: list[str], df: pd.DataFrame
 ) -> pd.DataFrame:
-    """
+    """Extract nested attributes from complex data structures into new DataFrame columns.
+
     This function, used in Online Tasks, is typically run on data exported from Arize.
-    It prepares the DataFrame by extracting relevant attributes from complex, deeply nested
-    data structures, such as those found in LLM outputs or JSON-like records. It helps extract
-    specific values from these nested structures by identifying the longest matching column name
-    in the DataFrame and recursively accessing the desired attribute path within each row.
-    This preprocessing step ensures that the extracted values are available as new columns,
-    allowing evaluators to process and assess these values effectively.
+    It prepares the DataFrame by extracting relevant attributes from complex, deeply
+    nested data structures, such as those found in LLM outputs or JSON-like records.
+    It helps extract specific values from these nested structures by identifying the
+    longest matching column name in the DataFrame and recursively accessing the desired
+    attribute path within each row. This preprocessing step ensures that the extracted
+    values are available as new columns, allowing evaluators to process and assess
+    these values effectively.
 
     For each attributes string in `attributes` (e.g. "attributes.llm.output_messages.0.message.content"),
     1) Find the largest prefix that is actually a column name in `df`. (e.g. "attributes.llm.output_messages")
@@ -37,13 +55,12 @@ def extract_nested_data_to_column(
     5) Log how many rows were dropped and, if zero rows remain, log a message indicating that
        there are no rows satisfying *all* of the queries.
     """
-
     # Make a copy so as not to alter the input df
     result_df = df.copy()
 
     # Keep track of which new columns we add. Each column name will match each user-inputted attribute
     # (e.g. "attributes.llm.output_messages.0.message.content")
-    new_cols: List[str] = []
+    new_cols: list[str] = []
 
     for attribute in attributes:
         parts = attribute.split(".")
@@ -58,7 +75,7 @@ def extract_nested_data_to_column(
                 prefix_len = i
 
         if prefix_col is None:
-            raise Exception("No such column found in DataFrame.")
+            raise ColumnNotFoundError(attribute)
 
         # 2) The remainder after the prefix
         remainder = ".".join(parts[prefix_len:])
@@ -68,13 +85,14 @@ def extract_nested_data_to_column(
             row: pd.Series,
             prefix_col: str = prefix_col,
             remainder: str = remainder,
-        ) -> Any:
+        ) -> object:
             val = row[prefix_col]
             try:
                 result = _introspect_arize_attribute(val, remainder)
-                return result if result is not None else np.nan
             except Exception:
                 return np.nan
+            else:
+                return result if result is not None else np.nan
 
         result_df[attribute] = result_df.apply(
             apply_introspect_arize_attribute, axis=1
@@ -101,9 +119,9 @@ def extract_nested_data_to_column(
     return result_df
 
 
-def _introspect_arize_attribute(value: Any, attribute: str) -> Any:
-    """
-    Recursively drill into `value` following the dot-delimited `attribute`.
+def _introspect_arize_attribute(value: object, attribute: str) -> object:
+    """Recursively drill into `value` following the dot-delimited `attribute`.
+
     Example:
         value: [{'message.role': 'assistant', 'message.content': 'The capital of China is Beijing.'}]
         attribute: "0.message.content"
@@ -124,8 +142,8 @@ def _introspect_arize_attribute(value: Any, attribute: str) -> Any:
 
 
 def _introspect_arize_attribute_parts(
-    current_value: Any, attribute_parts_unprocessed: List[str]
-) -> Any:
+    current_value: object, attribute_parts_unprocessed: list[str]
+) -> object:
     # If no more parts, we return whatever we have
     if not attribute_parts_unprocessed:
         return current_value
@@ -148,10 +166,9 @@ def _introspect_arize_attribute_parts(
 
 
 def _parse_value(
-    current_value: Any, attribute_parts_unprocessed: List[str]
-) -> Tuple[Any, int]:
-    """
-    Attempt to parse out the next value from `current_value` using the earliest parts:
+    current_value: object, attribute_parts_unprocessed: list[str]
+) -> tuple[object, int]:
+    """Attempt to parse out the next value from `current_value` using the earliest parts.
 
     1) If `attribute_parts_unprocessed[0]` is an integer index and `current_value` is a list/tuple,
        index into it.
@@ -164,7 +181,6 @@ def _parse_value(
       - parsed_value: the found value or None if not found
       - num_parts_processed: how many parts were processed (1 or more)
     """
-
     if not attribute_parts_unprocessed:
         return (None, 0)
 
@@ -182,35 +198,31 @@ def _parse_value(
         if isinstance(current_value, (list, tuple)):
             if 0 <= idx < len(current_value):
                 return (current_value[idx], num_parts_processed)
-            else:
-                return (None, num_parts_processed)
-        else:
             return (None, num_parts_processed)
+        return (None, num_parts_processed)
 
     # 2) Try dict approach
     if isinstance(current_value, dict):
         # a) direct match
         if key in current_value:
             return (current_value[key], num_parts_processed)
-        else:
-            # b) try combining multiple parts to handle dotted key
-            for num_parts_processed in range(
-                1, len(attribute_parts_unprocessed)
-            ):
-                key += "." + attribute_parts_unprocessed[num_parts_processed]
-                if key in current_value:
-                    return (
-                        current_value[key],
-                        num_parts_processed + 1,
-                    )
-            return (None, num_parts_processed)
+        # b) try combining multiple parts to handle dotted key
+        for num_parts_processed in range(1, len(attribute_parts_unprocessed)):
+            key += "." + attribute_parts_unprocessed[num_parts_processed]
+            if key in current_value:
+                return (
+                    current_value[key],
+                    num_parts_processed + 1,
+                )
+        return (None, num_parts_processed)
 
     # If we get here, we couldn't handle it (not a list or dict or mismatch)
     return (None, num_parts_processed)
 
 
-def _ensure_deserialized(val: Any) -> Any:
-    """
+def _ensure_deserialized(val: object) -> object:
+    """Ensure value is deserialized from numpy array or JSON string.
+
     1) If `val` is a numpy array, convert to a Python list.
     2) If `val` is a string, attempt to parse as JSON.
     3) Otherwise return as-is.

arize/utils/openinference_conversion.py

@@ -1,11 +1,24 @@
+"""OpenInference data conversion utilities for column transformations."""
+
 import json
+import logging
 
 import pandas as pd
 
 from arize.constants.openinference import OPEN_INFERENCE_JSON_STR_TYPES
 
+logger = logging.getLogger(__name__)
+
 
 def convert_datetime_columns_to_int(df: pd.DataFrame) -> pd.DataFrame:
+    """Convert datetime columns in a DataFrame to milliseconds since epoch.
+
+    Args:
+        df: The pandas DataFrame to convert.
+
+    Returns:
+        The DataFrame with datetime columns converted to integers.
+    """
     for col in df.select_dtypes(
         include=["datetime64[ns]", "datetime64[ns, UTC]"]
     ):
@@ -14,6 +27,14 @@ def convert_datetime_columns_to_int(df: pd.DataFrame) -> pd.DataFrame:
 
 
 def convert_boolean_columns_to_str(df: pd.DataFrame) -> pd.DataFrame:
+    """Convert boolean columns in a DataFrame to string type.
+
+    Args:
+        df: The pandas DataFrame to convert.
+
+    Returns:
+        The DataFrame with boolean columns converted to strings.
+    """
     for col in df.columns:
         if df[col].dtype == "bool":
             df[col] = df[col].astype("string")
@@ -21,33 +42,51 @@ def convert_boolean_columns_to_str(df: pd.DataFrame) -> pd.DataFrame:
 
 
 def convert_default_columns_to_json_str(df: pd.DataFrame) -> pd.DataFrame:
+    """Convert dictionary values in specific columns to JSON strings.
+
+    Args:
+        df: The pandas DataFrame to convert.
+
+    Returns:
+        The DataFrame with dictionaries in eligible columns converted to JSON strings.
+    """
     for col in df.columns:
         if _should_convert_json(col):
             try:
                 df[col] = df[col].apply(
                     lambda x: json.dumps(x) if isinstance(x, dict) else x
                 )
-            except Exception:
+            except Exception as e:
+                logger.debug(
+                    f"Failed to convert column '{col}' to JSON string: {e}"
+                )
                 continue
     return df
 
 
 def convert_json_str_to_dict(df: pd.DataFrame) -> pd.DataFrame:
+    """Convert JSON string values in specific columns to Python dictionaries.
+
+    Args:
+        df: The pandas DataFrame to convert.
+
+    Returns:
+        The DataFrame with JSON strings in eligible columns converted to dictionaries.
+    """
     for col in df.columns:
         if _should_convert_json(col):
             try:
                 df[col] = df[col].apply(
                     lambda x: json.loads(x) if isinstance(x, str) else x
                 )
-            except Exception:
+            except Exception as e:
+                logger.debug(f"Failed to parse column '{col}' as JSON: {e}")
                 continue
     return df
 
 
 def _should_convert_json(col_name: str) -> bool:
-    """
-    Check if a column should be converted to/from a JSON string/PythonDictionary.
-    """
+    """Check if a column should be converted to/from a JSON string/PythonDictionary."""
     is_eval_metadata = col_name.startswith("eval.") and col_name.endswith(
         ".metadata"
     )

arize/utils/proto.py

@@ -1,3 +1,5 @@
+"""Protocol buffer schema utilities for tracing data."""
+
 # type: ignore[pb2]
 from arize._generated.protocol.rec import public_pb2 as pb2
 
@@ -5,6 +7,14 @@ from arize._generated.protocol.rec import public_pb2 as pb2
 def get_pb_schema_tracing(
     project_name: str,
 ) -> pb2.Schema:
+    """Create a protobuf schema for LLM tracing data.
+
+    Args:
+        project_name: The name of the project/model.
+
+    Returns:
+        A configured pb2.Schema object for tracing environment.
+    """
     s = pb2.Schema()
     s.constants.model_id = project_name
     s.constants.environment = pb2.Schema.Environment.TRACING

arize/utils/size.py

@@ -1,13 +1,15 @@
+"""Size calculation utilities for payloads and data structures."""
+
 import sys
-from typing import Any, Dict, List
+from typing import Any
 
 import pandas as pd
 
 
-def get_payload_size_mb(payload: List[Dict[str, Any]] | pd.DataFrame) -> float:
+def get_payload_size_mb(payload: list[dict[str, Any]] | pd.DataFrame) -> float:
     """Return approximate size of payload in MB."""
     if isinstance(payload, pd.DataFrame):
-        # memory_usage(deep=True) sums all columns’ memory footprint
+        # memory_usage(deep=True) sums all columns' memory footprint
         size_bytes = payload.memory_usage(deep=True).sum()
     elif isinstance(payload, list):
         # sys.getsizeof() gives shallow size; sum all element sizes for rough total

arize/version.py

@@ -1 +1,3 @@
-__version__ = "8.0.0a21"
+"""Version information for the Arize SDK."""
+
+__version__ = "8.0.0a23"

{arize-8.0.0a21.dist-info → arize-8.0.0a23.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize
-Version: 8.0.0a21
+Version: 8.0.0a23
 Summary: A helper library to interact with Arize AI APIs
 Project-URL: Homepage, https://arize.com
 Project-URL: Documentation, https://docs.arize.com/arize
@@ -9,8 +9,9 @@ Project-URL: Source, https://github.com/Arize-ai/client_python
 Project-URL: Changelog, https://github.com/Arize-ai/client_python/blob/main/CHANGELOG.md
 Author-email: Arize AI <support@arize.com>
 Maintainer-email: Arize AI <support@arize.com>
-License: BSD
-License-File: LICENSE.md
+License: Apache-2.0
+License-File: LICENSE
+License-File: NOTICE
 Keywords: Arize,Evaluations,Explainability,LLM,Monitoring,Observability,Tracing
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers