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

arize/_flight/client.py

@@ -5,37 +5,34 @@ import base64
 import logging
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple
+from typing import TYPE_CHECKING
 
 from google.protobuf import json_format
 from pyarrow import flight
 
 from arize._flight.types import FlightRequestType
-from arize._generated.protocol.flight import ingest_pb2 as flight_ing_pb2
-from arize._generated.protocol.flight.ingest_pb2 import (
-    PostExperimentDataResponse,
-    WriteSpanAnnotationResponse,
-    WriteSpanAttributesMetadataResponse,
-    WriteSpanEvaluationResponse,
-)
-from arize.config import get_python_version
+from arize._generated.protocol.flight import flight_pb2
+from arize.config import PYTHON_VERSION
 from arize.logging import log_a_list
 from arize.utils.openinference_conversion import convert_json_str_to_dict
 from arize.utils.proto import get_pb_schema_tracing
 from arize.version import __version__
 
 if TYPE_CHECKING:
+    import types
+    from collections.abc import Iterable
+
     import pandas as pd
     import pyarrow as pa
 
 
-BytesPair = Tuple[bytes, bytes]
-Headers = List[BytesPair]
+BytesPair = tuple[bytes, bytes]
+Headers = list[BytesPair]
 FlightPostArrowFileResponse = (
-    WriteSpanEvaluationResponse
-    | WriteSpanAnnotationResponse
-    | WriteSpanAttributesMetadataResponse
-    | PostExperimentDataResponse
+    flight_pb2.WriteSpanEvaluationResponse
+    | flight_pb2.WriteSpanAnnotationResponse
+    | flight_pb2.WriteSpanAttributesMetadataResponse
+    | flight_pb2.PostExperimentDataResponse
 )
 
 logger = logging.getLogger(__name__)
@@ -51,6 +48,17 @@ class FlightActionKey(Enum):
 
 @dataclass(frozen=True)
 class ArizeFlightClient:
+    """Low-level Apache Arrow Flight client for Arize data operations.
+
+    This client provides the underlying Flight RPC transport for uploading and
+    downloading data to/from Arize using Apache Arrow format. It's used internally
+    by public-facing clients (SpansClient, ExperimentsClient, DatasetsClient, etc.)
+    and is not intended for direct use by end users.
+
+    The client manages connections, authentication, and protocol buffer serialization
+    for Flight operations.
+    """
+
     api_key: str = field(repr=False)
     host: str
     port: int
@@ -71,7 +79,7 @@ class ArizeFlightClient:
             (b"origin", b"arize-logging-client"),
             (b"auth-token-bin", str(self.api_key).encode("utf-8")),
             (b"sdk-language", b"python"),
-            (b"language-version", get_python_version().encode("utf-8")),
+            (b"language-version", PYTHON_VERSION.encode("utf-8")),
             (b"sdk-version", __version__.encode("utf-8")),
         ]
 
@@ -82,6 +90,7 @@ class ArizeFlightClient:
     # ---------- Connection management ----------
 
     def _ensure_client(self) -> flight.FlightClient:
+        """Lazily initialize and return the underlying Flight client connection."""
         client = object.__getattribute__(self, "_client")
         if client is not None:
             return client
@@ -99,6 +108,7 @@ class ArizeFlightClient:
         return new_client
 
     def close(self) -> None:
+        """Close the Flight client connection and clean up resources."""
         client = object.__getattribute__(self, "_client")
         if client is not None:
             client.close()
@@ -107,34 +117,58 @@ class ArizeFlightClient:
     # ---------- Context manager ----------
 
     def __enter__(self) -> ArizeFlightClient:
+        """Context manager entry point. Ensures the Flight client is initialized."""
         self._ensure_client()
         return self
 
-    def __exit__(self, exc_type, exc_val, _) -> None:
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        _: types.TracebackType | None,
+    ) -> None:
+        """Context manager exit point. Closes the Flight client connection."""
         if exc_type:
             logger.error(f"An exception occurred: {exc_val}")
         self.close()
 
     # ---------- methods simple passthrough wrappers ----------
 
-    def get_flight_info(self, *args: Any, **kwargs: Any):
+    def get_flight_info(self, *args: object, **kwargs: object) -> object:
+        """Get flight information. Passthrough to underlying Flight client with auth options."""
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.get_flight_info(*args, **kwargs)
 
-    def do_get(self, *args: Any, **kwargs: Any) -> flight.FlightStreamReader:
+    def do_get(
+        self, *args: object, **kwargs: object
+    ) -> flight.FlightStreamReader:
+        """Retrieve data stream via Flight DoGet.
+
+        Passthrough to underlying Flight client with auth options.
+        """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.do_get(*args, **kwargs)
 
     def do_put(
-        self, *args: Any, **kwargs: Any
+        self, *args: object, **kwargs: object
     ) -> [flight.FlightStreamWriter, flight.FlightMetadataReader]:
+        """Upload data stream via Flight DoPut.
+
+        Passthrough to underlying Flight client with auth options.
+        """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.do_put(*args, **kwargs)
 
-    def do_action(self, *args: Any, **kwargs: Any) -> Iterable[flight.Result]:
+    def do_action(
+        self, *args: object, **kwargs: object
+    ) -> Iterable[flight.Result]:
+        """Execute an action via Flight DoAction.
+
+        Passthrough to underlying Flight client with auth options.
+        """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.do_action(*args, **kwargs)
@@ -150,6 +184,26 @@ class ArizeFlightClient:
         dataset_id: str | None = None,
         experiment_name: str | None = None,
     ) -> FlightPostArrowFileResponse:
+        """Upload an Arrow table via Flight DoPut.
+
+        Args:
+            space_id: Space ID for the request.
+            request_type: Type of Flight request (EVALUATION, ANNOTATION, METADATA, or
+                LOG_EXPERIMENT_DATA).
+            pa_table: PyArrow Table containing the data to upload.
+            project_name: Project/model name (required for tracing request types).
+            dataset_id: Dataset ID (required for experiment request types).
+            experiment_name: Experiment name (required for experiment request types).
+
+        Returns:
+            FlightPostArrowFileResponse containing the server response, including counts
+            of processed records and any errors.
+
+        Raises:
+            ValueError: If request_type is unsupported or required parameters are missing
+                for the request type.
+            RuntimeError: If the Flight request fails.
+        """
         pa_schema = pa_table.schema
         if request_type in (
             FlightRequestType.EVALUATION,
@@ -187,31 +241,29 @@ class ArizeFlightClient:
                 if flight_response is None:
                     return None
 
-                res = None
-                match request_type:
-                    case FlightRequestType.EVALUATION:
-                        res = WriteSpanEvaluationResponse()
-                        res.ParseFromString(flight_response.to_pybytes())
-                    case FlightRequestType.ANNOTATION:
-                        res = WriteSpanAnnotationResponse()
-                        res.ParseFromString(flight_response.to_pybytes())
-                    case FlightRequestType.METADATA:
-                        res = WriteSpanAttributesMetadataResponse()
-                        res.ParseFromString(flight_response.to_pybytes())
-                    case FlightRequestType.LOG_EXPERIMENT_DATA:
-                        res = PostExperimentDataResponse()
-                        res.ParseFromString(flight_response.to_pybytes())
-                    case _:
-                        raise ValueError(
-                            f"Unsupported request_type: {request_type}"
-                        )
-                return res
         except Exception as e:
-            logger.exception(f"Error logging arrow table to Arize: {e}")
+            logger.exception("Error logging arrow table to Arize")
             raise RuntimeError(
                 f"Error logging arrow table to Arize: {e}"
             ) from e
 
+        match request_type:
+            case FlightRequestType.EVALUATION:
+                res = flight_pb2.WriteSpanEvaluationResponse()
+                res.ParseFromString(flight_response.to_pybytes())
+            case FlightRequestType.ANNOTATION:
+                res = flight_pb2.WriteSpanAnnotationResponse()
+                res.ParseFromString(flight_response.to_pybytes())
+            case FlightRequestType.METADATA:
+                res = flight_pb2.WriteSpanAttributesMetadataResponse()
+                res.ParseFromString(flight_response.to_pybytes())
+            case FlightRequestType.LOG_EXPERIMENT_DATA:
+                res = flight_pb2.PostExperimentDataResponse()
+                res.ParseFromString(flight_response.to_pybytes())
+            case _:
+                raise ValueError(f"Unsupported request_type: {request_type}")
+        return res
+
     # ---------- dataset methods ----------
 
     def create_dataset(
@@ -220,11 +272,24 @@ class ArizeFlightClient:
         dataset_name: str,
         pa_table: pa.Table,
     ) -> str:
-        doput_request = flight_ing_pb2.DoPutRequest(
-            create_dataset=flight_ing_pb2.CreateDatasetRequest(
+        """Create a new dataset via Flight DoPut.
+
+        Args:
+            space_id: Space ID where the dataset will be created.
+            dataset_name: Name for the new dataset.
+            pa_table: PyArrow Table containing the dataset examples.
+
+        Returns:
+            The created dataset ID as a string.
+
+        Raises:
+            RuntimeError: If the Flight request fails.
+        """
+        doput_request = flight_pb2.DoPutRequest(
+            create_dataset=flight_pb2.CreateDatasetRequest(
                 space_id=space_id,
                 dataset_name=dataset_name,
-                dataset_type=flight_ing_pb2.GENERATIVE,
+                dataset_type=flight_pb2.GENERATIVE,
             )
         )
         descriptor = flight.FlightDescriptor.for_command(
@@ -244,14 +309,11 @@ class ArizeFlightClient:
                 if flight_response is None:
                     return None
 
-                res = None
-                res = flight_ing_pb2.CreateDatasetResponse()
+                res = flight_pb2.CreateDatasetResponse()
                 res.ParseFromString(flight_response.to_pybytes())
-                if res:
-                    return str(res.dataset_id)
-                return res
+                return str(res.dataset_id) if res else None
         except Exception as e:
-            logger.exception(f"Error logging arrow table to Arize: {e}")
+            logger.exception("Error logging arrow table to Arize")
             raise RuntimeError(
                 f"Error logging arrow table to Arize: {e}"
             ) from e
@@ -262,10 +324,25 @@ class ArizeFlightClient:
         dataset_id: str,
         dataset_version_id: str | None = None,
     ) -> pd.DataFrame:
+        """Retrieve dataset examples via Flight DoGet.
+
+        Args:
+            space_id: Space ID containing the dataset.
+            dataset_id: Dataset ID to retrieve examples from.
+            dataset_version_id: Optional specific version ID. If None, retrieves the
+                latest version.
+
+        Returns:
+            A pandas DataFrame containing the dataset examples with JSON string columns
+            converted to dict objects.
+
+        Raises:
+            RuntimeError: If the Flight request fails.
+        """
         # TODO(Kiko): Space ID should not be needed,
         # should work on server tech debt to remove this
-        doget_request = flight_ing_pb2.DoGetRequest(
-            get_dataset=flight_ing_pb2.GetDatasetRequest(
+        doget_request = flight_pb2.DoGetRequest(
+            get_dataset=flight_pb2.GetDatasetRequest(
                 space_id=space_id,
                 dataset_id=dataset_id,
                 dataset_version=dataset_version_id,
@@ -278,8 +355,7 @@ class ArizeFlightClient:
             reader = self.do_get(descriptor, options=self.call_options)
             # read all data into pandas dataframe
             df = reader.read_all().to_pandas()
-            df = convert_json_str_to_dict(df)
-            return df
+            return convert_json_str_to_dict(df)
         except Exception as e:
             logger.exception(f"Failed to get dataset id={dataset_id}")
             raise RuntimeError(f"Failed to get dataset id={dataset_id}") from e
@@ -291,10 +367,23 @@ class ArizeFlightClient:
         space_id: str,
         experiment_id: str,
     ) -> pd.DataFrame:
+        """Retrieve experiment runs via Flight DoGet.
+
+        Args:
+            space_id: Space ID containing the experiment.
+            experiment_id: Experiment ID to retrieve runs from.
+
+        Returns:
+            A pandas DataFrame containing the experiment runs with JSON string columns
+            converted to dict objects.
+
+        Raises:
+            RuntimeError: If the Flight request fails.
+        """
         # TODO(Kiko): Space ID should not be needed,
         # should work on server tech debt to remove this
-        doget_request = flight_ing_pb2.DoGetRequest(
-            get_experiment=flight_ing_pb2.GetExperimentRequest(
+        doget_request = flight_pb2.DoGetRequest(
+            get_experiment=flight_pb2.GetExperimentRequest(
                 space_id=space_id,
                 experiment_id=experiment_id,
             )
@@ -306,8 +395,7 @@ class ArizeFlightClient:
             reader = self.do_get(descriptor, options=self.call_options)
             # read all data into pandas dataframe
             df = reader.read_all().to_pandas()
-            df = convert_json_str_to_dict(df)
-            return df
+            return convert_json_str_to_dict(df)
         except Exception as e:
             logger.exception(f"Failed to get experiment id={experiment_id}")
             raise RuntimeError(
@@ -319,9 +407,26 @@ class ArizeFlightClient:
         space_id: str,
         dataset_id: str,
         experiment_name: str,
-    ) -> Tuple[str, str] | None:
-        request = flight_ing_pb2.DoActionRequest(
-            create_experiment_db_entry=flight_ing_pb2.CreateExperimentDBEntryRequest(
+    ) -> tuple[str, str] | None:
+        """Initialize a new experiment via Flight DoAction.
+
+        Creates database entries for a new experiment and allocates a trace model for
+        capturing experiment traces.
+
+        Args:
+            space_id: Space ID where the experiment will be created.
+            dataset_id: Dataset ID the experiment will run against.
+            experiment_name: Name for the new experiment.
+
+        Returns:
+            A tuple of (experiment_id, trace_model_name) if successful, None if the
+            action response is empty.
+
+        Raises:
+            RuntimeError: If the Flight action fails.
+        """
+        request = flight_pb2.DoActionRequest(
+            create_experiment_db_entry=flight_pb2.CreateExperimentDBEntryRequest(
                 space_id=space_id,
                 dataset_id=dataset_id,
                 experiment_name=experiment_name,
@@ -342,7 +447,7 @@ class ArizeFlightClient:
         res = next(response, None)
         if res is None:
             return None
-        resp_pb = flight_ing_pb2.CreateExperimentDBEntryResponse()
+        resp_pb = flight_pb2.CreateExperimentDBEntryResponse()
         resp_pb.ParseFromString(res.body.to_pybytes())
         return (
             resp_pb.experiment_id,
@@ -351,8 +456,20 @@ class ArizeFlightClient:
 
 
 def append_to_pyarrow_metadata(
-    pa_schema: pa.Schema, new_metadata: Dict[str, Any]
-):
+    pa_schema: pa.Schema, new_metadata: dict[str, object]
+) -> object:
+    """Append metadata to a PyArrow schema without overwriting existing metadata.
+
+    Args:
+        pa_schema: The PyArrow schema to add metadata to.
+        new_metadata: Dictionary of metadata key-value pairs to append.
+
+    Returns:
+        A new PyArrow schema with the combined metadata.
+
+    Raises:
+        KeyError: If any keys in new_metadata already exist in the schema's metadata.
+    """
     # Ensure metadata is handled correctly, even if initially None.
     metadata = pa_schema.metadata
     if metadata is None:
@@ -377,9 +494,23 @@ def _get_pb_flight_doput_request(
     model_id: str | None = None,
     dataset_id: str | None = None,
     experiment_name: str | None = None,
-) -> flight_ing_pb2.DoPutRequest:
-    """Return a DoPutRequest for the given request_type."""
-
+) -> flight_pb2.DoPutRequest:
+    """Construct a Flight DoPut protocol buffer request for the given request type.
+
+    Args:
+        space_id: Space ID for the request.
+        request_type: FlightRequestType enum value.
+        model_id: Optional model/project ID (required for tracing request types).
+        dataset_id: Optional dataset ID (required for experiment request types).
+        experiment_name: Optional experiment name (required for experiment request types).
+
+    Returns:
+        A DoPutRequest protocol buffer configured for the specified request type.
+
+    Raises:
+        ValueError: If the request_type is unsupported or required parameters are
+            missing for the given request_type.
+    """
     common_args = {"space_id": space_id}
 
     if model_id:
@@ -392,21 +523,21 @@ def _get_pb_flight_doput_request(
     if model_id:
         # model-based request types
         match request_type:
-            case FlightRequestType.LOG_EXPERIMENT_DATA:
-                return flight_ing_pb2.DoPutRequest(
-                    write_span_evaluation_request=flight_ing_pb2.WriteSpanEvaluationRequest(
+            case FlightRequestType.EVALUATION:
+                return flight_pb2.DoPutRequest(
+                    write_span_evaluation_request=flight_pb2.WriteSpanEvaluationRequest(
                         **common_args
                     )
                 )
             case FlightRequestType.ANNOTATION:
-                return flight_ing_pb2.DoPutRequest(
-                    write_span_annotation_request=flight_ing_pb2.WriteSpanAnnotationRequest(
+                return flight_pb2.DoPutRequest(
+                    write_span_annotation_request=flight_pb2.WriteSpanAnnotationRequest(
                         **common_args
                     )
                 )
             case FlightRequestType.METADATA:
-                return flight_ing_pb2.DoPutRequest(
-                    write_span_attributes_metadata_request=flight_ing_pb2.WriteSpanAttributesMetadataRequest(
+                return flight_pb2.DoPutRequest(
+                    write_span_attributes_metadata_request=flight_pb2.WriteSpanAttributesMetadataRequest(
                         **common_args
                     )
                 )
@@ -417,8 +548,8 @@ def _get_pb_flight_doput_request(
         # dataset-based request types
         match request_type:
             case FlightRequestType.LOG_EXPERIMENT_DATA:
-                return flight_ing_pb2.DoPutRequest(
-                    post_experiment_data=flight_ing_pb2.PostExperimentDataRequest(
+                return flight_pb2.DoPutRequest(
+                    post_experiment_data=flight_pb2.PostExperimentDataRequest(
                         **common_args
                     )
                 )

arize/_generated/api_client/__init__.py

@@ -5,9 +5,9 @@
 """
     Arize REST API
 
-    API specification for the backend data server. The API is hosted globally at https://api.arize.com/v2 or in your own environment. You can access the OpenAPI spec for this API at https://api.arize.com/v2/spec.yaml 
+    API specification for the backend data server. The API is hosted globally at https://api.arize.com/v2 or in your own environment. 
 
-    The version of the OpenAPI document: 0.0.1
+    The version of the OpenAPI document: 2.0.0
     Generated by OpenAPI Generator (https://openapi-generator.tech)
 
     Do not edit the class manually.
@@ -20,6 +20,7 @@ __version__ = "1.0.0"
 __all__ = [
     "DatasetsApi",
     "ExperimentsApi",
+    "ProjectsApi",
     "ApiResponse",
     "ApiClient",
     "Configuration",
@@ -30,20 +31,32 @@ __all__ = [
     "ApiAttributeError",
     "ApiException",
     "Dataset",
+    "DatasetExample",
+    "DatasetExampleUpdate",
     "DatasetVersion",
     "DatasetsCreateRequest",
+    "DatasetsExamplesInsertRequest",
+    "DatasetsExamplesList200Response",
+    "DatasetsExamplesUpdateRequest",
     "DatasetsList200Response",
-    "DatasetsListExamples200Response",
-    "Error",
     "Experiment",
+    "ExperimentRun",
+    "ExperimentRunCreate",
     "ExperimentsCreateRequest",
     "ExperimentsList200Response",
     "ExperimentsRunsList200Response",
+    "PaginationMetadata",
+    "PrimitiveValue",
+    "Problem",
+    "Project",
+    "ProjectsCreateRequest",
+    "ProjectsList200Response",
 ]
 
 # import apis into sdk package
 from arize._generated.api_client.api.datasets_api import DatasetsApi as DatasetsApi
 from arize._generated.api_client.api.experiments_api import ExperimentsApi as ExperimentsApi
+from arize._generated.api_client.api.projects_api import ProjectsApi as ProjectsApi
 
 # import ApiClient
 from arize._generated.api_client.api_response import ApiResponse as ApiResponse
@@ -58,13 +71,24 @@ from arize._generated.api_client.exceptions import ApiException as ApiException
 
 # import models into sdk package
 from arize._generated.api_client.models.dataset import Dataset as Dataset
+from arize._generated.api_client.models.dataset_example import DatasetExample as DatasetExample
+from arize._generated.api_client.models.dataset_example_update import DatasetExampleUpdate as DatasetExampleUpdate
 from arize._generated.api_client.models.dataset_version import DatasetVersion as DatasetVersion
 from arize._generated.api_client.models.datasets_create_request import DatasetsCreateRequest as DatasetsCreateRequest
+from arize._generated.api_client.models.datasets_examples_insert_request import DatasetsExamplesInsertRequest as DatasetsExamplesInsertRequest
+from arize._generated.api_client.models.datasets_examples_list200_response import DatasetsExamplesList200Response as DatasetsExamplesList200Response
+from arize._generated.api_client.models.datasets_examples_update_request import DatasetsExamplesUpdateRequest as DatasetsExamplesUpdateRequest
 from arize._generated.api_client.models.datasets_list200_response import DatasetsList200Response as DatasetsList200Response
-from arize._generated.api_client.models.datasets_list_examples200_response import DatasetsListExamples200Response as DatasetsListExamples200Response
-from arize._generated.api_client.models.error import Error as Error
 from arize._generated.api_client.models.experiment import Experiment as Experiment
+from arize._generated.api_client.models.experiment_run import ExperimentRun as ExperimentRun
+from arize._generated.api_client.models.experiment_run_create import ExperimentRunCreate as ExperimentRunCreate
 from arize._generated.api_client.models.experiments_create_request import ExperimentsCreateRequest as ExperimentsCreateRequest
 from arize._generated.api_client.models.experiments_list200_response import ExperimentsList200Response as ExperimentsList200Response
 from arize._generated.api_client.models.experiments_runs_list200_response import ExperimentsRunsList200Response as ExperimentsRunsList200Response
+from arize._generated.api_client.models.pagination_metadata import PaginationMetadata as PaginationMetadata
+from arize._generated.api_client.models.primitive_value import PrimitiveValue as PrimitiveValue
+from arize._generated.api_client.models.problem import Problem as Problem
+from arize._generated.api_client.models.project import Project as Project
+from arize._generated.api_client.models.projects_create_request import ProjectsCreateRequest as ProjectsCreateRequest
+from arize._generated.api_client.models.projects_list200_response import ProjectsList200Response as ProjectsList200Response
 

arize/_generated/api_client/api/__init__.py

@@ -3,4 +3,5 @@
 # import apis into api package
 from arize._generated.api_client.api.datasets_api import DatasetsApi
 from arize._generated.api_client.api.experiments_api import ExperimentsApi
+from arize._generated.api_client.api.projects_api import ProjectsApi