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

arize/__init__.py

@@ -1,9 +1,12 @@
+"""Arize SDK for model observability and LLM tracing."""
+
 import logging
 from collections.abc import Mapping
 
 from arize._generated.api_client import models
 from arize.client import ArizeClient
 from arize.config import SDKConfiguration
+from arize.regions import Region
 
 # Attach a NullHandler by default in the top-level package
 # so that if no configuration is installed, nothing explodes.
@@ -14,23 +17,27 @@ try:
     from .logging import auto_configure_from_env
 
     auto_configure_from_env()
-except Exception:
-    # Never let logging config crash imports
+except Exception:  # noqa: S110
+    # Intentionally silent: logging configuration is optional and should never
+    # prevent SDK initialization. Users can configure logging explicitly if needed.
     pass
 
-__all__ = ["ArizeClient", "SDKConfiguration"]
+__all__ = [
+    "ArizeClient",
+    "Region",
+    "SDKConfiguration",
+]
 
 
-def make_to_df(field_name: str):
+def make_to_df(field_name: str) -> object:
     def to_df(
-        self,
+        self: object,
         by_alias: bool = False,
         exclude_none: str | bool = False,
         json_normalize: bool = False,
         convert_dtypes: bool = True,
-    ):
-        """
-        Convert a list of objects to a pandas DataFrame.
+    ) -> object:
+        """Convert a list of objects to a pandas DataFrame.
 
         Behavior:
           - If an item is a Pydantic v2 model, use `.model_dump(by_alias=...)`.
@@ -85,6 +92,7 @@ def make_to_df(field_name: str):
 
 
 models.DatasetsList200Response.to_df = make_to_df("datasets")  # type: ignore[attr-defined]
-models.DatasetsListExamples200Response.to_df = make_to_df("examples")  # type: ignore[attr-defined]
+models.DatasetsExamplesList200Response.to_df = make_to_df("examples")  # type: ignore[attr-defined]
 models.ExperimentsList200Response.to_df = make_to_df("experiments")  # type: ignore[attr-defined]
 models.ExperimentsRunsList200Response.to_df = make_to_df("experiment_runs")  # type: ignore[attr-defined]
+models.ProjectsList200Response.to_df = make_to_df("projects")  # type: ignore[attr-defined]

arize/_exporter/client.py

@@ -2,7 +2,6 @@
 import logging
 from dataclasses import dataclass
 from datetime import datetime
-from typing import List, Tuple
 
 import pandas as pd
 import pyarrow.parquet as pq
@@ -16,7 +15,7 @@ from arize._exporter.validation import (
     validate_input_type,
     validate_start_end_time,
 )
-from arize._generated.protocol.flight import export_pb2
+from arize._generated.protocol.flight import flight_pb2
 from arize.logging import CtxAdapter
 from arize.types import Environments, SimilaritySearchParams
 from arize.utils.dataframe import reset_dataframe_index
@@ -33,19 +32,20 @@ class ArizeExportClient:
         space_id: str,
         model_id: str,
         environment: Environments,
-        start_time: str | datetime,
-        end_time: str | datetime,
+        start_time: datetime,
+        end_time: datetime,
         where: str = "",
-        columns: List | None = None,
+        columns: list | None = None,
         similarity_search_params: SimilaritySearchParams | None = None,
         model_version: str = "",
         batch_id: str = "",
         include_actuals: bool = False,
         stream_chunk_size: int | None = None,
-    ):
-        """
-        Exports data of a specific model in the Arize platform to a pandas dataframe for a defined
-        time interval and model environment, optionally by model version and/or batch id.
+    ) -> object:
+        """Exports data of a specific model in the Arize platform to a pandas dataframe.
+
+        The export covers a defined time interval and model environment, and can
+        optionally be filtered by model version and/or batch id.
 
         Args:
             space_id (str): The id for the space where to export models from, can be retrieved from
@@ -104,14 +104,14 @@ class ArizeExportClient:
             return pd.DataFrame()
         progress_bar = self._get_progress_bar(num_recs)
         list_of_df = []
-        while True:
-            try:
+        try:
+            while True:
                 flight_batch = stream_reader.read_chunk()
                 batch_df = flight_batch.data.to_pandas()
                 list_of_df.append(batch_df)
                 progress_bar.update(batch_df.shape[0])
-            except StopIteration:
-                break
+        except StopIteration:
+            pass
         progress_bar.close()
         df = pd.concat(list_of_df)
         null_columns = df.columns[df.isnull().all()]
@@ -139,16 +139,17 @@ class ArizeExportClient:
         start_time: datetime,
         end_time: datetime,
         where: str = "",
-        columns: List | None = None,
+        columns: list | None = None,
         similarity_search_params: SimilaritySearchParams | None = None,
         model_version: str = "",
         batch_id: str = "",
         include_actuals: bool = False,
         stream_chunk_size: int | None = None,
     ) -> None:
-        """
-        Exports data of a specific model in the Arize platform to a parquet file for a defined time
-        interval and model environment, optionally by model version and/or batch id.
+        """Exports data of a specific model in the Arize platform to a parquet file.
+
+        The export covers a defined time interval and model environment, and can
+        optionally be filtered by model version and/or batch id.
 
         Args:
             path (str): path to the file to store exported data. File must be in parquet format and
@@ -208,17 +209,17 @@ class ArizeExportClient:
             stream_chunk_size=stream_chunk_size,
         )
         if stream_reader is None:
-            return None
+            return
         progress_bar = self._get_progress_bar(num_recs)
         with pq.ParquetWriter(path, schema=stream_reader.schema) as writer:
-            while True:
-                try:
+            try:
+                while True:
                     flight_batch = stream_reader.read_chunk()
                     record_batch = flight_batch.data
                     writer.write_batch(record_batch)
                     progress_bar.update(record_batch.num_rows)
-                except StopIteration:
-                    break
+            except StopIteration:
+                pass
         progress_bar.close()
 
     def _get_stream_reader(
@@ -233,9 +234,9 @@ class ArizeExportClient:
         batch_id: str = "",
         where: str = "",
         similarity_search_params: SimilaritySearchParams | None = None,
-        columns: List | None = None,
+        columns: list | None = None,
         stream_chunk_size: int | None = None,
-    ) -> Tuple[flight.FlightStreamReader, int]:
+    ) -> tuple[flight.FlightStreamReader | None, int]:
         # Bind common context for this operation
         log = CtxAdapter(
             logger,
@@ -273,7 +274,7 @@ class ArizeExportClient:
         validate_start_end_time(start_time, end_time)
 
         # Create query descriptor
-        query_descriptor = export_pb2.RecordQueryDescriptor(
+        query_descriptor = flight_pb2.RecordQueryDescriptor(
             space_id=space_id,
             model_id=model_id,
             environment=environment.name,
@@ -289,9 +290,11 @@ class ArizeExportClient:
                 else None
             ),
             projected_columns=columns if columns else [],
-            stream_chunk_size=Int64Value(value=stream_chunk_size)
-            if stream_chunk_size is not None
-            else None,
+            stream_chunk_size=(
+                Int64Value(value=stream_chunk_size)
+                if stream_chunk_size is not None
+                else None
+            ),
         )
 
         try:
@@ -306,17 +309,24 @@ class ArizeExportClient:
                 logger.warning("Query returns no data")
                 return None, 0
             logger.debug("Ticket: %s", flight_info.endpoints[0].ticket)
-
-            # Retrieve the result set as flight stream reader
-            reader = self.flight_client.do_get(flight_info.endpoints[0].ticket)
-            return reader, flight_info.total_records
         except Exception as e:
             msg = f"Error getting flight info or do_get: {e}"
-            logger.error(msg)
+            logger.exception(msg)
             raise RuntimeError(msg) from e
+        # Retrieve the result set as flight stream reader
+        reader = self.flight_client.do_get(flight_info.endpoints[0].ticket)
+        return reader, flight_info.total_records
 
     @staticmethod
-    def _get_progress_bar(num_recs):
+    def _get_progress_bar(num_recs: int) -> tqdm:
+        """Create a progress bar for export operations.
+
+        Args:
+            num_recs: Total number of records to export.
+
+        Returns:
+            A tqdm progress bar configured for data export display.
+        """
         return tqdm(
             total=num_recs,
             desc=f"  exporting {num_recs} rows",
@@ -329,8 +339,17 @@ class ArizeExportClient:
 
 def _get_pb_similarity_search_params(
     similarity_params: SimilaritySearchParams,
-) -> export_pb2.SimilaritySearchParams:
-    proto_params = export_pb2.SimilaritySearchParams()
+) -> flight_pb2.SimilaritySearchParams:
+    """Convert SimilaritySearchParams to protocol buffer format.
+
+    Args:
+        similarity_params: Similarity search parameters containing search column name,
+            threshold, and reference examples.
+
+    Returns:
+        A protocol buffer SimilaritySearchParams object for Flight requests.
+    """
+    proto_params = flight_pb2.SimilaritySearchParams()
     proto_params.search_column_name = similarity_params.search_column_name
     proto_params.threshold = similarity_params.threshold
     for ref in similarity_params.references:

arize/_exporter/parsers/tracing_data_parser.py

@@ -1,6 +1,5 @@
 import json
 import logging
-from typing import List
 
 import numpy as np
 import pandas as pd
@@ -28,12 +27,26 @@ logger = logging.getLogger(__name__)
 # but the resulting error messages provide clarity on what the effect
 # of the error is on the data; It should not prevent a user from continuing to use the data
 class OtelTracingDataTransformer:
+    def _apply_column_transformation(
+        self, df: pd.DataFrame, col_name: str, transform_func: object
+    ) -> str | None:
+        """Apply a transformation to a column and return error message if it fails."""
+        try:
+            df[col_name] = df[col_name].apply(transform_func)
+        except Exception as e:
+            return (
+                f"Unable to transform json string data to a Python dict in column '{col_name}'; "
+                f"May encounter issues when importing data back into Arize; Error: {e}"
+            )
+        else:
+            return None
+
     def transform(self, df: pd.DataFrame) -> pd.DataFrame:
-        errors: List[str] = []
+        errors: list[str] = []
 
         # Convert list of json serializable strings columns to list of dictionaries for more
         # conveinent data processing in Python
-        list_of_json_string_column_names: List[str] = [
+        list_of_json_string_column_names: list[str] = [
             col.name
             for col in [
                 SPAN_ATTRIBUTES_LLM_INPUT_MESSAGES_COL,
@@ -45,17 +58,13 @@ class OtelTracingDataTransformer:
             if col.name in df.columns
         ]
         for col_name in list_of_json_string_column_names:
-            try:
-                df[col_name] = df[col_name].apply(
-                    self._transform_value_to_list_of_dict
-                )
-            except Exception as e:
-                errors.append(
-                    f"Unable to transform json string data to a Python dict in column '{col_name}'; "
-                    f"May encounter issues when importing data back into Arize; Error: {e}"
-                )
-
-        json_string_column_names: List[str] = [
+            error = self._apply_column_transformation(
+                df, col_name, self._transform_value_to_list_of_dict
+            )
+            if error:
+                errors.append(error)
+
+        json_string_column_names: list[str] = [
             col.name
             for col in [
                 SPAN_ATTRIBUTES_LLM_PROMPT_TEMPLATE_VARIABLES_COL,
@@ -64,16 +73,14 @@ class OtelTracingDataTransformer:
             if col.name in df.columns
         ]
         for col_name in json_string_column_names:
-            try:
-                df[col_name] = df[col_name].apply(self._transform_json_to_dict)
-            except Exception as e:
-                errors.append(
-                    f"Unable to transform json string data to a Python dict in column '{col_name}'; "
-                    f"May encounter issues when importing data back into Arize; Error: {e}"
-                )
+            error = self._apply_column_transformation(
+                df, col_name, self._transform_json_to_dict
+            )
+            if error:
+                errors.append(error)
 
         # Clean json string columns since empty strings are equivalent here to None but are not valid json
-        dirty_string_column_names: List[str] = [
+        dirty_string_column_names: list[str] = [
             col.name
             for col in [
                 SPAN_ATTRIBUTES_LLM_INVOCATION_PARAMETERS_COL,
@@ -85,7 +92,7 @@ class OtelTracingDataTransformer:
             df[col_name] = df[col_name].apply(self._clean_json_string)
 
         # Convert timestamp columns to datetime objects
-        timestamp_column_names: List[str] = [
+        timestamp_column_names: list[str] = [
             col.name
             for col in [
                 SPAN_START_TIME_COL,
@@ -103,7 +110,9 @@ class OtelTracingDataTransformer:
 
         return df
 
-    def _transform_value_to_list_of_dict(self, value):
+    def _transform_value_to_list_of_dict(
+        self, value: object
+    ) -> list[object] | None:
         if value is None:
             return None
 
@@ -113,10 +122,11 @@ class OtelTracingDataTransformer:
                 for i in value
                 if self._is_non_empty_string(i)
             ]
-        elif self._is_non_empty_string(value):
+        if self._is_non_empty_string(value):
             return [self._deserialize_json_string_to_dict(value)]
+        return None
 
-    def _transform_json_to_dict(self, value):
+    def _transform_json_to_dict(self, value: object) -> object | None:
         if value is None:
             return None
 
@@ -126,20 +136,21 @@ class OtelTracingDataTransformer:
         if isinstance(value, str) and value == "":
             # transform empty string to None
             return None
+        return None
 
-    def _is_non_empty_string(self, value):
+    def _is_non_empty_string(self, value: object) -> bool:
         return isinstance(value, str) and value != ""
 
-    def _deserialize_json_string_to_dict(self, value: str):
+    def _deserialize_json_string_to_dict(self, value: str) -> object:
         try:
             return json.loads(value)
         except json.JSONDecodeError as e:
             raise ValueError(f"Invalid JSON string: {value}") from e
 
-    def _clean_json_string(self, value):
+    def _clean_json_string(self, value: object) -> object | None:
         return value if self._is_non_empty_string(value) else None
 
-    def _convert_timestamp_to_datetime(self, value):
+    def _convert_timestamp_to_datetime(self, value: object) -> object:
         return (
             pd.Timestamp(value, unit="ns")
             if value and isinstance(value, (int, float, np.int64))

arize/_exporter/validation.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from datetime import datetime
 
 
 def validate_input_type(
-    input: Any,
+    input: object,
     input_name: str,
     input_type: type,
     allow_none: bool = False,
@@ -28,7 +28,7 @@ def validate_input_type(
 
 
 def validate_input_value(
-    input: Any,
+    input: object,
     input_name: str,
     choices: tuple,
 ) -> None: