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

arize/spans/validation/spans/value_validation.py

@@ -1,10 +1,13 @@
+"""Value validation logic for span data."""
+
 from __future__ import annotations
 
 from itertools import chain
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 
 from arize.constants import spans as tracing_constants
 from arize.constants.ml import MAX_EMBEDDING_DIMENSIONALITY
+from arize.ml.types import StatusCodes
 from arize.spans import columns as tracing_cols
 from arize.spans.validation.common import value_validation
 from arize.spans.validation.common.errors import (
@@ -13,7 +16,7 @@ from arize.spans.validation.common.errors import (
     InvalidEventValueInColumn,
     InvalidLLMMessageValueInColumn,
 )
-from arize.types import StatusCodes, is_dict_of, is_json_str
+from arize.utils.types import is_dict_of, is_json_str
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -23,7 +26,18 @@ if TYPE_CHECKING:
 
 def check_span_root_field_values(
     dataframe: pd.DataFrame,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
+    """Validates root-level span field values for proper format and constraints.
+
+    Checks span ID, trace ID, parent span ID, name, status code, status message,
+    timestamps, and events for conformance to specification limits.
+
+    Args:
+        dataframe: The DataFrame containing span data.
+
+    Returns:
+        List of validation errors found in root span fields.
+    """
     return list(
         chain(
             value_validation.check_string_column_value_length(
@@ -77,7 +91,18 @@ def check_span_root_field_values(
 
 def check_span_attributes_values(
     dataframe: pd.DataFrame,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
+    """Validates span attribute values for proper format and constraints.
+
+    Checks all span attributes including LLM parameters, embeddings, documents,
+    tools, and other metadata fields for conformance to specification limits.
+
+    Args:
+        dataframe: The DataFrame containing span data.
+
+    Returns:
+        List of validation errors found in span attributes.
+    """
     return list(
         chain(
             value_validation.check_string_column_value_length(
@@ -242,7 +267,17 @@ def check_span_attributes_values(
 
 def check_event_column_value(
     df: pd.DataFrame,
-) -> List[InvalidEventValueInColumn]:
+) -> list[InvalidEventValueInColumn]:
+    """Validates span event column values for proper format and length constraints.
+
+    Checks event names for length limits and attributes for proper dictionary structure.
+
+    Args:
+        df: The DataFrame containing span events.
+
+    Returns:
+        List of validation errors found in event column values.
+    """
     col_name = tracing_cols.SPAN_EVENTS_COL.name
     if col_name not in df.columns:
         return []
@@ -284,7 +319,18 @@ def check_event_column_value(
 
 def check_embeddings_column_value(
     df: pd.DataFrame,
-) -> List[InvalidEmbeddingValueInColumn]:
+) -> list[InvalidEmbeddingValueInColumn]:
+    """Validates embedding column values for proper vector dimensions and text length.
+
+    Checks that embedding vectors are within dimensionality limits and text
+    values don't exceed maximum length.
+
+    Args:
+        df: The DataFrame containing embedding data.
+
+    Returns:
+        List of validation errors found in embedding column values.
+    """
     col_name = tracing_cols.SPAN_ATTRIBUTES_EMBEDDING_EMBEDDINGS_COL.name
     if col_name not in df.columns:
         return []
@@ -332,7 +378,19 @@ def check_embeddings_column_value(
 def check_LLM_IO_messages_column_value(
     df: pd.DataFrame,
     col_name: str,
-) -> List[InvalidLLMMessageValueInColumn]:
+) -> list[InvalidLLMMessageValueInColumn]:
+    """Validates LLM input/output message column values for proper format and length.
+
+    Checks message role, content, and tool calls for conformance to length limits
+    and proper JSON formatting.
+
+    Args:
+        df: The DataFrame containing LLM messages.
+        col_name: Name of the message column to validate.
+
+    Returns:
+        List of validation errors found in message column values.
+    """
     if col_name not in df.columns:
         return []
 
@@ -407,7 +465,19 @@ def check_LLM_IO_messages_column_value(
 def check_documents_column_value(
     df: pd.DataFrame,
     col_name: str,
-) -> List[InvalidDocumentValueInColumn]:
+) -> list[InvalidDocumentValueInColumn]:
+    """Validates document column values for proper format and length constraints.
+
+    Checks document ID, content, and metadata for conformance to length limits
+    and proper data type requirements.
+
+    Args:
+        df: The DataFrame containing documents.
+        col_name: Name of the document column to validate.
+
+    Returns:
+        List of validation errors found in document column values.
+    """
     if col_name not in df.columns:
         return []
 

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,13 +1,19 @@
+"""DataFrame manipulation and validation utilities."""
+
 import re
-from typing import List
 
 import pandas as pd
 
-from arize.types import BaseSchema
+from arize.ml.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,10 +119,10 @@ 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`.
-    Example:
+def _introspect_arize_attribute(value: object, attribute: str) -> object:
+    """Recursively drill into `value` following the dot-delimited `attribute`.
+
+    Examples:
         value: [{'message.role': 'assistant', 'message.content': 'The capital of China is Beijing.'}]
         attribute: "0.message.content"
         Returns: 'The capital of China is Beijing.'
@@ -114,7 +132,6 @@ def _introspect_arize_attribute(value: Any, attribute: str) -> Any:
       - Parses JSON strings
       - Converts NumPy arrays to lists
       - Allows dotted keys (e.g. "message.content") by combining parts
-
     """
     if not attribute:
         return value
@@ -124,8 +141,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 +165,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 +180,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)
 
@@ -179,38 +194,34 @@ def _parse_value(
     idx = _try_int(key)
     if idx is not None:
         # Must be a tuple or list (_ensure_deserialized() already casts numpy arrays to python lists)
-        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)
+        if isinstance(current_value, list | tuple) and 0 <= idx < len(
+            current_value
+        ):
+            return (current_value[idx], 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"
     )