arize 8.0.0a15__py3-none-any.whl → 8.0.0a17__py3-none-any.whl

arize/__init__.py

@@ -1,6 +1,7 @@
-# src/arize/__init__.py
 import logging
+from collections.abc import Mapping
 
+from arize._generated.api_client import models
 from arize.client import ArizeClient
 from arize.config import SDKConfiguration
 
@@ -18,3 +19,72 @@ except Exception:
     pass
 
 __all__ = ["ArizeClient", "SDKConfiguration"]
+
+
+def make_to_df(field_name: str):
+    def to_df(
+        self,
+        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.
+
+        Behavior:
+          - If an item is a Pydantic v2 model, use `.model_dump(by_alias=...)`.
+          - If an item is a mapping (dict-like), use it as-is.
+          - Otherwise, raise a ValueError (unsupported row type).
+
+        Parameters:
+          by_alias: Use field aliases when dumping Pydantic models.
+          exclude_none:
+            - False: keep Nones as-is
+            - "all": drop columns where *all* values are None/NaN
+            - "any": drop columns where *any* value is None/NaN
+            - True: alias for "all"
+          json_normalize: If True, flatten nested dicts via `pandas.json_normalize`.
+          convert_dtypes: If True, call `DataFrame.convert_dtypes()` at the end.
+
+        Returns:
+          pandas.DataFrame
+        """
+        import pandas as pd
+
+        items = getattr(self, field_name, []) or []
+
+        rows = []
+        for it in items:
+            if hasattr(it, "model_dump"):  # Pydantic v2 object
+                rows.append(it.model_dump(by_alias=by_alias))
+
+            elif isinstance(it, Mapping):  # Plain mapping
+                rows.append(it)
+            else:
+                raise ValueError(
+                    f"Cannot convert item of type {type(it)} to DataFrame row"
+                )
+
+        df = (
+            pd.json_normalize(rows, sep=".")
+            if json_normalize
+            else pd.DataFrame(rows)
+        )
+
+        # Drop None/NaN columns if requested
+        if exclude_none in ("any", "all", True):
+            drop_how = "all" if exclude_none is True else exclude_none
+            df.dropna(axis=1, how=drop_how, inplace=True)
+
+        if convert_dtypes:
+            df = df.convert_dtypes()
+        return df
+
+    return to_df
+
+
+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.ExperimentsList200Response.to_df = make_to_df("experiments")  # type: ignore[attr-defined]
+models.ExperimentsRunsList200Response.to_df = make_to_df("experiment_runs")  # type: ignore[attr-defined]

arize/_flight/client.py

@@ -4,7 +4,8 @@ from __future__ import annotations
 import base64
 import logging
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, List, Tuple
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple
 
 from google.protobuf import json_format
 from pyarrow import flight
@@ -12,35 +13,49 @@ 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.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 pandas as pd
     import pyarrow as pa
 
 
 BytesPair = Tuple[bytes, bytes]
 Headers = List[BytesPair]
-WriteSpanResponse = (
+FlightPostArrowFileResponse = (
     WriteSpanEvaluationResponse
     | WriteSpanAnnotationResponse
     | WriteSpanAttributesMetadataResponse
+    | PostExperimentDataResponse
 )
 
 logger = logging.getLogger(__name__)
 
 
+class FlightActionKey(Enum):
+    CREATE_EXPERIMENT_DB_ENTRY = "create_experiment_db_entry"
+    # GET_DATASET_VERSION = "get_dataset_version"
+    # LIST_DATASETS = "list_datasets"
+    # DELETE_DATASET = "delete_dataset"
+    # DELETE_EXPERIMENT = "delete_experiment"
+
+
 @dataclass(frozen=True)
 class ArizeFlightClient:
     api_key: str = field(repr=False)
     host: str
     port: int
     scheme: str
+    max_chunksize: int
     request_verify: bool
 
     # internal cache for the underlying FlightClient
@@ -107,17 +122,19 @@ class ArizeFlightClient:
         kwargs.setdefault("options", self.call_options)
         return client.get_flight_info(*args, **kwargs)
 
-    def do_get(self, *args: Any, **kwargs: Any):
+    def do_get(self, *args: Any, **kwargs: Any) -> flight.FlightStreamReader:
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.do_get(*args, **kwargs)
 
-    def do_put(self, *args: Any, **kwargs: Any):
+    def do_put(
+        self, *args: Any, **kwargs: Any
+    ) -> [flight.FlightStreamWriter, flight.FlightMetadataReader]:
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.do_put(*args, **kwargs)
 
-    def do_action(self, *args: Any, **kwargs: Any):
+    def do_action(self, *args: Any, **kwargs: Any) -> Iterable[flight.Result]:
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.do_action(*args, **kwargs)
@@ -127,20 +144,30 @@ class ArizeFlightClient:
     def log_arrow_table(
         self,
         space_id: str,
-        project_name: str,
         request_type: FlightRequestType,
         pa_table: pa.Table,
-    ) -> WriteSpanResponse:
-        proto_schema = get_pb_schema_tracing(project_name=project_name)
-        base64_schema = base64.b64encode(proto_schema.SerializeToString())
-        pa_schema = append_to_pyarrow_metadata(
-            pa_table.schema, {"arize-schema": base64_schema}
-        )
+        project_name: str | None = None,
+        dataset_id: str | None = None,
+        experiment_name: str | None = None,
+    ) -> FlightPostArrowFileResponse:
+        pa_schema = pa_table.schema
+        if request_type in (
+            FlightRequestType.EVALUATION,
+            FlightRequestType.ANNOTATION,
+            FlightRequestType.METADATA,
+        ):
+            proto_schema = get_pb_schema_tracing(project_name=project_name)
+            base64_schema = base64.b64encode(proto_schema.SerializeToString())
+            pa_schema = append_to_pyarrow_metadata(
+                pa_table.schema, {"arize-schema": base64_schema}
+            )
 
         doput_request = _get_pb_flight_doput_request(
             space_id=space_id,
-            model_id=project_name,
             request_type=request_type,
+            model_id=project_name,
+            dataset_id=dataset_id,
+            experiment_name=experiment_name,
         )
 
         descriptor = flight.FlightDescriptor.for_command(
@@ -152,7 +179,7 @@ class ArizeFlightClient:
             )
             with flight_writer:
                 # write table as stream to flight server
-                flight_writer.write_table(pa_table)
+                flight_writer.write_table(pa_table, self.max_chunksize)
                 # indicate that client has flushed all contents to stream
                 flight_writer.done_writing()
                 # read response from flight server
@@ -171,11 +198,13 @@ class ArizeFlightClient:
                     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}")
@@ -207,7 +236,7 @@ class ArizeFlightClient:
             )
             with flight_writer:
                 # write table as stream to flight server
-                flight_writer.write_table(pa_table)
+                flight_writer.write_table(pa_table, self.max_chunksize)
                 # indicate that client has flushed all contents to stream
                 flight_writer.done_writing()
                 # read response from flight server
@@ -227,6 +256,99 @@ class ArizeFlightClient:
                 f"Error logging arrow table to Arize: {e}"
             ) from e
 
+    def get_dataset_examples(
+        self,
+        space_id: str,
+        dataset_id: str,
+        dataset_version_id: str | None = None,
+    ) -> pd.DataFrame:
+        # 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(
+                space_id=space_id,
+                dataset_id=dataset_id,
+                dataset_version=dataset_version_id,
+            )
+        )
+        descriptor = flight.Ticket(
+            json_format.MessageToJson(doget_request).encode("utf-8")
+        )
+        try:
+            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
+        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
+
+    # ---------- experiment methods ----------
+
+    def get_experiment_runs(
+        self,
+        space_id: str,
+        experiment_id: str,
+    ) -> pd.DataFrame:
+        # 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(
+                space_id=space_id,
+                experiment_id=experiment_id,
+            )
+        )
+        descriptor = flight.Ticket(
+            json_format.MessageToJson(doget_request).encode("utf-8")
+        )
+        try:
+            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
+        except Exception as e:
+            logger.exception(f"Failed to get experiment id={experiment_id}")
+            raise RuntimeError(
+                f"Failed to get experiment id={experiment_id}"
+            ) from e
+
+    def init_experiment(
+        self,
+        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(
+                space_id=space_id,
+                dataset_id=dataset_id,
+                experiment_name=experiment_name,
+            )
+        )
+        action = flight.Action(
+            FlightActionKey.CREATE_EXPERIMENT_DB_ENTRY,
+            json_format.MessageToJson(request).encode("utf-8"),
+        )
+        try:
+            response = self.do_action(action, options=self.call_options)
+        except Exception as e:
+            logger.exception(f"Failed to init experiment {experiment_name}")
+            raise RuntimeError(
+                f"Failed to init experiment {experiment_name}"
+            ) from e
+
+        res = next(response, None)
+        if res is None:
+            return None
+        resp_pb = flight_ing_pb2.CreateExperimentDBEntryResponse()
+        resp_pb.ParseFromString(res.body.to_pybytes())
+        return (
+            resp_pb.experiment_id,
+            resp_pb.trace_model_name,
+        )
+
 
 def append_to_pyarrow_metadata(
     pa_schema: pa.Schema, new_metadata: Dict[str, Any]
@@ -250,34 +372,59 @@ def append_to_pyarrow_metadata(
 
 
 def _get_pb_flight_doput_request(
-    space_id,
-    model_id: str,
+    space_id: str,
     request_type: FlightRequestType,
+    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."""
-    common_args: dict[str, str] = {
-        "space_id": space_id,
-        "external_model_id": model_id,
-    }
-
-    match request_type:
-        case FlightRequestType.EVALUATION:
-            return flight_ing_pb2.DoPutRequest(
-                write_span_evaluation_request=flight_ing_pb2.WriteSpanEvaluationRequest(
-                    **common_args
+
+    common_args = {"space_id": space_id}
+
+    if model_id:
+        common_args["external_model_id"] = model_id
+    if dataset_id:
+        common_args["dataset_id"] = dataset_id
+    if experiment_name:
+        common_args["experiment_name"] = experiment_name
+
+    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(
+                        **common_args
+                    )
                 )
-            )
-        case FlightRequestType.ANNOTATION:
-            return flight_ing_pb2.DoPutRequest(
-                write_span_annotation_request=flight_ing_pb2.WriteSpanAnnotationRequest(
-                    **common_args
+            case FlightRequestType.ANNOTATION:
+                return flight_ing_pb2.DoPutRequest(
+                    write_span_annotation_request=flight_ing_pb2.WriteSpanAnnotationRequest(
+                        **common_args
+                    )
                 )
-            )
-        case FlightRequestType.METADATA:
-            return flight_ing_pb2.DoPutRequest(
-                write_span_attributes_metadata_request=flight_ing_pb2.WriteSpanAttributesMetadataRequest(
-                    **common_args
+            case FlightRequestType.METADATA:
+                return flight_ing_pb2.DoPutRequest(
+                    write_span_attributes_metadata_request=flight_ing_pb2.WriteSpanAttributesMetadataRequest(
+                        **common_args
+                    )
                 )
-            )
-        case _:
-            raise ValueError(f"Unsupported request_type: {request_type}")
+            case _:
+                raise ValueError(f"Unsupported request_type: {request_type}")
+
+    if dataset_id and experiment_name:
+        # dataset-based request types
+        match request_type:
+            case FlightRequestType.LOG_EXPERIMENT_DATA:
+                return flight_ing_pb2.DoPutRequest(
+                    post_experiment_data=flight_ing_pb2.PostExperimentDataRequest(
+                        **common_args
+                    )
+                )
+            case _:
+                raise ValueError(f"Unsupported request_type: {request_type}")
+
+    raise ValueError(
+        f"Unsupported combination: {request_type=} with provided arguments."
+    )

arize/_flight/types.py

@@ -5,3 +5,4 @@ class FlightRequestType(str, Enum):
     EVALUATION = "evaluation"
     ANNOTATION = "annotation"
     METADATA = "metadata"
+    LOG_EXPERIMENT_DATA = "log_experiment_data"

arize/_generated/api_client/__init__.py

@@ -5,7 +5,7 @@
 """
     Arize REST API
 
-    API specification for the backend data server. The API is hosted globally at https://app.arize.com/api/v1 or in your own environment. You can access the OpenAPI spec for this API at https://app.arize.com/api/v1/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. You can access the OpenAPI spec for this API at https://api.arize.com/v2/spec.yaml 
 
     The version of the OpenAPI document: 0.0.1
     Generated by OpenAPI Generator (https://openapi-generator.tech)
@@ -36,7 +36,9 @@ __all__ = [
     "DatasetsListExamples200Response",
     "Error",
     "Experiment",
+    "ExperimentsCreateRequest",
     "ExperimentsList200Response",
+    "ExperimentsRunsList200Response",
 ]
 
 # import apis into sdk package
@@ -62,5 +64,7 @@ from arize._generated.api_client.models.datasets_list200_response import Dataset
 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.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
 

arize/_generated/api_client/api/datasets_api.py

@@ -3,7 +3,7 @@
 """
     Arize REST API
 
-    API specification for the backend data server. The API is hosted globally at https://app.arize.com/api/v1 or in your own environment. You can access the OpenAPI spec for this API at https://app.arize.com/api/v1/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. You can access the OpenAPI spec for this API at https://api.arize.com/v2/spec.yaml 
 
     The version of the OpenAPI document: 0.0.1
     Generated by OpenAPI Generator (https://openapi-generator.tech)
@@ -309,7 +309,7 @@ class DatasetsApi:
 
         return self.api_client.param_serialize(
             method='POST',
-            resource_path='/api/v1/datasets',
+            resource_path='/v2/datasets',
             path_params=_path_params,
             query_params=_query_params,
             header_params=_header_params,
@@ -585,7 +585,7 @@ class DatasetsApi:
 
         return self.api_client.param_serialize(
             method='DELETE',
-            resource_path='/api/v1/datasets/{datasetId}',
+            resource_path='/v2/datasets/{datasetId}',
             path_params=_path_params,
             query_params=_query_params,
             header_params=_header_params,
@@ -878,7 +878,7 @@ class DatasetsApi:
 
         return self.api_client.param_serialize(
             method='GET',
-            resource_path='/api/v1/datasets/{datasetId}',
+            resource_path='/v2/datasets/{datasetId}',
             path_params=_path_params,
             query_params=_query_params,
             header_params=_header_params,
@@ -1170,7 +1170,7 @@ class DatasetsApi:
 
         return self.api_client.param_serialize(
             method='GET',
-            resource_path='/api/v1/datasets',
+            resource_path='/v2/datasets',
             path_params=_path_params,
             query_params=_query_params,
             header_params=_header_params,
@@ -1480,7 +1480,7 @@ class DatasetsApi:
 
         return self.api_client.param_serialize(
             method='GET',
-            resource_path='/api/v1/datasets/{datasetId}/examples',
+            resource_path='/v2/datasets/{datasetId}/examples',
             path_params=_path_params,
             query_params=_query_params,
             header_params=_header_params,