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

arize/spans/conversion.py

@@ -1,6 +1,8 @@
+"""Span data conversion utilities for transforming and normalizing span data."""
+
 import json
-from datetime import datetime
-from typing import Any, Dict, Iterable, List
+from collections.abc import Iterable
+from datetime import datetime, timezone
 
 import numpy as np
 import pandas as pd
@@ -10,53 +12,62 @@ from arize.spans.columns import SPAN_OPENINFERENCE_COLUMNS, SpanColumnDataType
 
 
 def convert_timestamps(df: pd.DataFrame, fmt: str = "") -> pd.DataFrame:
+    """Convert timestamp columns in a DataFrame to nanoseconds.
+
+    Args:
+        df: The pandas DataFrame containing timestamp columns.
+        fmt: Optional datetime format string for parsing string timestamps. Defaults to "".
+
+    Returns:
+        The DataFrame with timestamp columns converted to nanoseconds.
+
+    Raises:
+        KeyError: If required timestamp column is not found in DataFrame.
+    """
     for col in SPAN_OPENINFERENCE_COLUMNS:
         if col.data_type != SpanColumnDataType.TIMESTAMP:
             continue
+        if col.name not in df.columns:
+            raise KeyError(f"Column '{col.name}' not found in DataFrame")
         df[col.name] = df[col.name].apply(lambda dt: _datetime_to_ns(dt, fmt))
     return df
 
 
 def _datetime_to_ns(dt: object, fmt: str) -> int:
     if isinstance(dt, str):
-        try:
-            ts = int(datetime.timestamp(datetime.strptime(dt, fmt)) * 1e9)
-        except Exception as e:
-            # logger.error(
-            #     f"Error parsing string '{dt}' to timestamp in nanoseconds "
-            #     f"using the format '{fmt}': {e}"
-            # )
-            raise e
-        return ts
-    elif isinstance(dt, datetime):
-        try:
-            ts = int(datetime.timestamp(dt) * 1e9)
-        except Exception as e:
-            # logger.error(
-            #     f"Error converting datetime object to nanoseconds: {e}"
-            # )
-            raise e
-        return ts
-    elif isinstance(dt, pd.Timestamp):
+        return int(
+            datetime.strptime(dt, fmt).replace(tzinfo=timezone.utc).timestamp()
+            * 1e9
+        )
+    if isinstance(dt, datetime):
+        return int(datetime.timestamp(dt) * 1e9)
+    if isinstance(dt, pd.Timestamp):
         return int(dt.value)
-    elif isinstance(dt, pd.DatetimeIndex):
+    if isinstance(dt, pd.DatetimeIndex):
         # Only allow a single element; otherwise ambiguous for a scalar function
         if len(dt) != 1:
             raise TypeError(
                 f"Expected a single timestamp in DatetimeIndex, got length={len(dt)}"
             )
         return int(dt.to_numpy(dtype="datetime64[ns]").astype("int64")[0])
-    elif isinstance(dt, (int, float)):
+    if isinstance(dt, (int, float)):
         # Assume value already in nanoseconds,
         # validate timestamps in validate_values
         return int(dt)
-    else:
-        e = TypeError(f"Cannot convert type {type(dt)} to nanoseconds")
-        # logger.error(f"Error converting pandas Timestamp to nanoseconds: {e}")
-        raise e
+    e = TypeError(f"Cannot convert type {type(dt)} to nanoseconds")
+    # logger.error(f"Error converting pandas Timestamp to nanoseconds: {e}")
+    raise e
 
 
 def jsonify_dictionaries(df: pd.DataFrame) -> pd.DataFrame:
+    """Convert dictionary and list-of-dictionary columns to JSON strings.
+
+    Args:
+        df: The pandas DataFrame containing dictionary columns.
+
+    Returns:
+        The DataFrame with dictionary columns converted to JSON strings.
+    """
     # NOTE: numpy arrays are not json serializable. Hence, we assume the
     # embeddings come as lists, not arrays
     dict_cols = [
@@ -90,7 +101,15 @@ def jsonify_dictionaries(df: pd.DataFrame) -> pd.DataFrame:
 
 
 # Defines what is considered a missing value
-def isMissingValue(value: Any) -> bool:
+def is_missing_value(value: object) -> bool:
+    """Check if a value should be considered missing or invalid.
+
+    Args:
+        value: The value to check.
+
+    Returns:
+        True if the value is missing (NaN, infinity, or pandas NA), False otherwise.
+    """
     assumed_missing_values = (
         np.inf,
         -np.inf,
@@ -99,22 +118,19 @@ def isMissingValue(value: Any) -> bool:
 
 
 def _jsonify_list_of_dicts(
-    list_of_dicts: Iterable[Dict[str, Any]] | None,
-) -> List[str]:
-    if not isinstance(list_of_dicts, Iterable) and isMissingValue(
+    list_of_dicts: Iterable[dict[str, object]] | None,
+) -> list[str]:
+    if not isinstance(list_of_dicts, Iterable) and is_missing_value(
         list_of_dicts
     ):
         return []
-    list_of_json = []
-    for d in list_of_dicts:
-        list_of_json.append(_jsonify_dict(d))
-    return list_of_json
+    return [_jsonify_dict(d) for d in list_of_dicts]
 
 
-def _jsonify_dict(d: Dict[str, Any] | None) -> str | None:
+def _jsonify_dict(d: dict[str, object] | None) -> str | None:
     if d is None:
-        return
-    if isMissingValue(d):
+        return None
+    if is_missing_value(d):
         return None
     d = d.copy()  # avoid side effects
     for k, v in d.items():

arize/spans/validation/__init__.py

@@ -0,0 +1 @@
+"""Validation utilities for LLM tracing spans data."""

arize/spans/validation/annotations/__init__.py

@@ -0,0 +1 @@
+"""Annotation validation for LLM tracing spans."""

arize/spans/validation/annotations/annotations_validation.py

@@ -1,7 +1,9 @@
+"""Annotation validation orchestration for spans."""
+
 from __future__ import annotations
 
 from itertools import chain
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 
 from arize.spans.columns import SPAN_SPAN_ID_COL
 from arize.spans.validation.annotations import (
@@ -27,7 +29,7 @@ if TYPE_CHECKING:
 def validate_argument_types(
     annotations_dataframe: pd.DataFrame,
     project_name: str,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
     """Validates argument types for log_annotations."""
     checks = chain(
         common_arg_validation.check_field_convertible_to_str(project_name),
@@ -40,7 +42,7 @@ def validate_argument_types(
 
 def validate_dataframe_form(
     annotations_dataframe: pd.DataFrame,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
     """Validates the form/structure of the annotation dataframe."""
     # Call annotation-specific function (to be created)
     df_validation.log_info_dataframe_extra_column_names(annotations_dataframe)
@@ -64,7 +66,7 @@ def validate_dataframe_form(
 def validate_values(
     annotations_dataframe: pd.DataFrame,
     project_name: str,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
     """Validates the values within the annotation dataframe."""
     checks = chain(
         # Common checks remain the same

arize/spans/validation/annotations/dataframe_form_validation.py

@@ -1,8 +1,10 @@
+"""DataFrame form validation for span annotations."""
+
 from __future__ import annotations
 
 import logging
 import re
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 
 import pandas as pd
 
@@ -19,7 +21,7 @@ from arize.spans.columns import (
     ANNOTATION_UPDATED_BY_SUFFIX,
     SPAN_SPAN_ID_COL,
 )
-from arize.spans.conversion import isMissingValue
+from arize.spans.conversion import is_missing_value
 from arize.spans.validation.common.errors import (
     InvalidAnnotationColumnFormat,
     InvalidDataFrameColumnContentTypes,
@@ -36,7 +38,7 @@ def log_info_dataframe_extra_column_names(
 ) -> None:
     """Logs columns that don't match expected annotation or context patterns."""
     if df is None:
-        return None
+        return
     # Check against annotation pattern, span id, and note column
     irrelevant_columns = [
         col
@@ -56,12 +58,12 @@ def log_info_dataframe_extra_column_names(
             "- annotation.<your-annotation-name>.score"
             f"An optional '{ANNOTATION_NOTES_COLUMN_NAME}' column can also be included."
         )
-    return None
+    return
 
 
 def check_invalid_annotation_column_names(
     df: pd.DataFrame,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
     """Checks for columns that start with 'annotation.' but don't match the expected pattern."""
     errors = []
 
@@ -86,7 +88,7 @@ def check_invalid_annotation_column_names(
 
 def check_dataframe_column_content_type(
     df: pd.DataFrame,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
     """Checks that columns matching annotation patterns have the correct data types."""
     wrong_labels_cols = []
     wrong_scores_cols = []
@@ -128,14 +130,14 @@ def check_dataframe_column_content_type(
         # Check annotation label column type (string or missing)
         elif annotation_label_re.match(column):
             if not all(
-                isinstance(value, str) or isMissingValue(value)
+                isinstance(value, str) or is_missing_value(value)
                 for value in df[column]
             ):
                 wrong_labels_cols.append(column)
         # Check annotation score column type (numeric or missing)
         elif annotation_score_re.match(column):
             if not all(
-                isinstance(value, (int, float)) or isMissingValue(value)
+                isinstance(value, (int, float)) or is_missing_value(value)
                 for value in df[column]
             ):
                 wrong_scores_cols.append(column)
@@ -144,21 +146,21 @@ def check_dataframe_column_content_type(
             if not all(
                 # Note: After formatting, this column holds list<string> (JSON), not just string.
                 # We rely on later schema inference/validation. Keep basic check for now.
-                isinstance(value, list) or isMissingValue(value)
+                isinstance(value, list) or is_missing_value(value)
                 for value in df[column]
             ):
                 wrong_notes_cols.append(column)
         # Check annotation updated_by column type (string or missing)
         elif annotation_updated_by_re.match(column):
             if not all(
-                isinstance(value, str) or isMissingValue(value)
+                isinstance(value, str) or is_missing_value(value)
                 for value in df[column]
             ):
                 wrong_updated_by_cols.append(column)
         # Check annotation updated_at column type (numeric or missing)
         elif annotation_updated_at_re.match(column) and not all(
             # Allow int, float (e.g., Unix timestamp millis)
-            isinstance(value, (int, float)) or isMissingValue(value)
+            isinstance(value, (int, float)) or is_missing_value(value)
             for value in df[column]
         ):
             wrong_updated_at_cols.append(column)

arize/spans/validation/annotations/value_validation.py

@@ -1,10 +1,12 @@
+"""Value validation logic for span annotation data."""
+
 from __future__ import annotations
 
 import logging
 import re
-from datetime import datetime
+from datetime import datetime, timezone
 from itertools import chain
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 
 from arize.constants.spans import (
     ANNOTATION_LABEL_MAX_STR_LENGTH,
@@ -41,33 +43,55 @@ logger = logging.getLogger(__name__)
 
 
 class InvalidAnnotationTimestamp(ValidationError):
+    """Raised when annotation timestamp is invalid or out of acceptable range."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Annotation_Timestamp"
 
     def __init__(self, timestamp_col_name: str, error_type: str) -> None:
+        """Initialize the exception with timestamp validation context.
+
+        Args:
+            timestamp_col_name: Name of the annotation timestamp column.
+            error_type: Type of timestamp error (e.g., 'future').
+        """
         self.timestamp_col_name = timestamp_col_name
         self.error_type = error_type
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         if self.error_type == "future":
             return (
                 f"At least one timestamp in the annotation column '{self.timestamp_col_name}' "
                 f"is in the future. Annotation timestamps cannot be in the future."
             )
-        elif self.error_type == "non_positive":
+        if self.error_type == "non_positive":
             return (
                 f"At least one timestamp in the annotation column '{self.timestamp_col_name}' "
                 f"is zero or negative. Annotation timestamps must be positive values."
             )
-        else:
-            return f"Invalid timestamp in annotation column '{self.timestamp_col_name}'."
+        return f"Invalid timestamp in annotation column '{self.timestamp_col_name}'."
 
 
 def check_annotation_updated_at_timestamp(
     df: pd.DataFrame,
     col_name: str,
     is_required: bool,
-) -> List[InvalidMissingValueInColumn | InvalidAnnotationTimestamp]:
+) -> list[InvalidMissingValueInColumn | InvalidAnnotationTimestamp]:
+    """Validates annotation timestamp values for validity and acceptable ranges.
+
+    Checks that timestamp values are positive, not in the future, and satisfy
+    required constraints if specified.
+
+    Args:
+        df: DataFrame containing the annotation timestamp column.
+        col_name: Name of the timestamp column to validate.
+        is_required: Whether the column must have non-null values in all rows.
+
+    Returns:
+        List of validation errors found (empty if valid).
+    """
     # This check expects that timestamps have previously been converted to milliseconds
     if col_name not in df.columns:
         return []
@@ -83,7 +107,7 @@ def check_annotation_updated_at_timestamp(
     if df[col_name].isnull().all():
         return errors
 
-    now_ms = datetime.now().timestamp() * 1000
+    now_ms = datetime.now(tz=timezone.utc).timestamp() * 1000
 
     if df[col_name].max() > now_ms:
         logger.warning(f"Detected future timestamp in column '{col_name}'.")
@@ -105,7 +129,7 @@ def check_annotation_updated_at_timestamp(
 
 def check_annotation_cols(
     dataframe: pd.DataFrame,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
     """Checks value length and validity for columns matching annotation patterns."""
     checks = []
     for col in dataframe.columns:
@@ -150,7 +174,7 @@ def check_annotation_cols(
 
 def check_annotation_columns_null_values(
     dataframe: pd.DataFrame,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
     """Checks that for a given annotation name, at least one of label or score is non-null per row."""
     invalid_annotation_names = []
     annotation_names = set()
@@ -190,7 +214,7 @@ def check_annotation_columns_null_values(
                 invalid_annotation_names.append(ann_name)
 
     # Use set to report each name only once
-    unique_invalid_names = sorted(list(set(invalid_annotation_names)))
+    unique_invalid_names = sorted(set(invalid_annotation_names))
     if unique_invalid_names:
         return [
             InvalidNullAnnotationLabelAndScore(
@@ -202,7 +226,7 @@ def check_annotation_columns_null_values(
 
 def check_annotation_notes_column(
     dataframe: pd.DataFrame,
-) -> List[ValidationError]:
+) -> list[ValidationError]:
     """Checks the value length for the optional annotation.notes column (raw string)."""
     col_name = ANNOTATION_NOTES_COLUMN_NAME
     if col_name in dataframe.columns:

arize/spans/validation/common/__init__.py

@@ -0,0 +1 @@
+"""Common validation utilities shared across spans validation."""

arize/spans/validation/common/argument_validation.py

@@ -1,4 +1,4 @@
-from typing import Any, List
+"""Common argument validation utilities for spans."""
 
 import pandas as pd
 
@@ -7,9 +7,18 @@ from arize.spans.validation.common.errors import InvalidTypeArgument
 
 
 def check_field_convertible_to_str(
-    project_name: Any,
-    model_version: Any = None,
-) -> List[InvalidFieldTypeConversion]:
+    project_name: object,
+    model_version: object = None,
+) -> list[InvalidFieldTypeConversion]:
+    """Validates that field arguments can be converted to strings.
+
+    Args:
+        project_name: The project name value to validate for string conversion.
+        model_version: Optional model version value to validate for string conversion.
+
+    Returns:
+        List of validation errors for fields that cannot be converted to strings.
+    """
     wrong_fields = []
     if project_name is not None and not isinstance(project_name, str):
         try:
@@ -28,8 +37,16 @@ def check_field_convertible_to_str(
 
 
 def check_dataframe_type(
-    dataframe,
-) -> List[InvalidTypeArgument]:
+    dataframe: object,
+) -> list[InvalidTypeArgument]:
+    """Validates that the provided argument is a pandas DataFrame.
+
+    Args:
+        dataframe: The object to validate as a pandas DataFrame.
+
+    Returns:
+        List of validation errors if not a DataFrame (empty if valid).
+    """
     if not isinstance(dataframe, pd.DataFrame):
         return [
             InvalidTypeArgument(
@@ -42,8 +59,16 @@ def check_dataframe_type(
 
 
 def check_datetime_format_type(
-    dt_fmt: Any,
-) -> List[InvalidTypeArgument]:
+    dt_fmt: object,
+) -> list[InvalidTypeArgument]:
+    """Validates that the datetime format argument is a string.
+
+    Args:
+        dt_fmt: The datetime format value to validate.
+
+    Returns:
+        List of validation errors if not a string (empty if valid).
+    """
     if not isinstance(dt_fmt, str):
         return [
             InvalidTypeArgument(

arize/spans/validation/common/dataframe_form_validation.py

@@ -1,6 +1,8 @@
+"""Common DataFrame form validation for spans."""
+
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 
 from arize.exceptions.base import InvalidDataFrameIndex
 from arize.spans.validation.common.errors import (
@@ -14,7 +16,15 @@ if TYPE_CHECKING:
 
 def check_dataframe_index(
     dataframe: pd.DataFrame,
-) -> List[InvalidDataFrameIndex]:
+) -> list[InvalidDataFrameIndex]:
+    """Validates that the DataFrame has a default integer index.
+
+    Args:
+        dataframe: The DataFrame to validate.
+
+    Returns:
+        List of validation errors if index is not default (empty if valid).
+    """
     if (dataframe.index != dataframe.reset_index(drop=True).index).any():
         return [InvalidDataFrameIndex()]
     return []
@@ -22,13 +32,21 @@ def check_dataframe_index(
 
 def check_dataframe_required_column_set(
     df: pd.DataFrame,
-    required_columns: List[str],
-) -> List[InvalidDataFrameMissingColumns]:
+    required_columns: list[str],
+) -> list[InvalidDataFrameMissingColumns]:
+    """Validates that the DataFrame contains all required columns.
+
+    Args:
+        df: The DataFrame to validate.
+        required_columns: List of column names that must be present.
+
+    Returns:
+        List of validation errors for missing columns (empty if valid).
+    """
     existing_columns = set(df.columns)
-    missing_cols = []
-    for col in required_columns:
-        if col not in existing_columns:
-            missing_cols.append(col)
+    missing_cols = [
+        col for col in required_columns if col not in existing_columns
+    ]
 
     if missing_cols:
         return [InvalidDataFrameMissingColumns(missing_cols=missing_cols)]
@@ -37,7 +55,15 @@ def check_dataframe_required_column_set(
 
 def check_dataframe_for_duplicate_columns(
     df: pd.DataFrame,
-) -> List[InvalidDataFrameDuplicateColumns]:
+) -> list[InvalidDataFrameDuplicateColumns]:
+    """Validates that the DataFrame has no duplicate column names.
+
+    Args:
+        df: The DataFrame to validate.
+
+    Returns:
+        List of validation errors if duplicate columns exist (empty if valid).
+    """
     # Get the duplicated column names from the dataframe
     duplicate_columns = df.columns[df.columns.duplicated()]
     if not duplicate_columns.empty: