arize 8.0.0a14__tar.gz → 8.0.0a16__tar.gz

{arize-8.0.0a14 → arize-8.0.0a16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize
-Version: 8.0.0a14
+Version: 8.0.0a16
 Summary: A helper library to interact with Arize AI APIs
 Project-URL: Homepage, https://arize.com
 Project-URL: Documentation, https://docs.arize.com/arize
@@ -35,7 +35,9 @@ Requires-Dist: tokenizers<1,>=0.13; extra == 'auto-embeddings'
 Requires-Dist: torch<3,>=1.13; extra == 'auto-embeddings'
 Requires-Dist: transformers<5,>=4.25; extra == 'auto-embeddings'
 Provides-Extra: datasets-experiments
+Requires-Dist: numpy>=2.0.0; extra == 'datasets-experiments'
 Requires-Dist: pydantic; extra == 'datasets-experiments'
+Requires-Dist: wrapt<2.0.0,>=1.0.0; extra == 'datasets-experiments'
 Provides-Extra: dev
 Requires-Dist: pytest==8.4.2; extra == 'dev'
 Requires-Dist: ruff==0.13.2; extra == 'dev'

{arize-8.0.0a14 → arize-8.0.0a16}/pyproject.toml

@@ -38,7 +38,7 @@ dependencies = [
     "lazy-imports",
     # "requests_futures==1.0.0",
     # "googleapis_common_protos>=1.51.0,<2",
-    # "protobuf>=4.21.0,<6",
+    # "protobuf>=4.21.0,<7",
     # "pyarrow>=0.15.0",
     # "tqdm>=4.60.0,<5",
     # "pydantic>=2.0.0,<3",
@@ -72,6 +72,15 @@ ml-batch = [
 ]
 datasets-experiments = [
     "pydantic",
+    "numpy>=2.0.0",
+    "wrapt>=1.0.0,<2.0.0",
+    # "openinference-semantic-conventions>=0.1.21, <1",
+    # "opentelemetry-exporter-otlp-proto-common>=1.38.0",
+    # "opentelemetry-exporter-otlp-proto-grpc>=1.38.0",
+    # "opentelemetry-sdk>=1.38.0",
+    #
+    # "opentelemetry-api>=1.38.0",
+    # "opentelemetry-proto>=1.38.0",
 ]
 mimic-explainer = [
     "interpret-community[mimic]>=0.22.0,<1",

{arize-8.0.0a14/src/arize/datasets → arize-8.0.0a16/src/arize}/__init__.py

@@ -1,6 +1,24 @@
+import logging
 from collections.abc import Mapping
 
 from arize._generated.api_client import models
+from arize.client import ArizeClient
+from arize.config import SDKConfiguration
+
+# Attach a NullHandler by default in the top-level package
+# so that if no configuration is installed, nothing explodes.
+logging.getLogger("arize").addHandler(logging.NullHandler())
+
+# Opt-in env-based logging
+try:
+    from .logging import auto_configure_from_env
+
+    auto_configure_from_env()
+except Exception:
+    # Never let logging config crash imports
+    pass
+
+__all__ = ["ArizeClient", "SDKConfiguration"]
 
 
 def make_to_df(field_name: str):
@@ -68,3 +86,4 @@ 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.ExperimentsList200Response.to_df = make_to_df("experiments")  # type: ignore[attr-defined]

{arize-8.0.0a14 → arize-8.0.0a16}/src/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,12 +13,15 @@ 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:
@@ -26,21 +30,31 @@ if TYPE_CHECKING:
 
 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 +121,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 +143,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 +178,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,15 +197,19 @@ 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}")
-                    raise RuntimeError(f"Error logging arrow table to Arize: {e}") from e
+            raise RuntimeError(
+                f"Error logging arrow table to Arize: {e}"
+            ) from e
 
     # ---------- dataset methods ----------
 
@@ -205,7 +235,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
@@ -221,7 +251,72 @@ class ArizeFlightClient:
                 return res
         except Exception as e:
             logger.exception(f"Error logging arrow table to Arize: {e}")
-            raise RuntimeError(f"Error logging arrow table to Arize: {e}") from e
+            raise RuntimeError(
+                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,
+    ):
+        # 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
+
+    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(
@@ -246,34 +341,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-8.0.0a14 → arize-8.0.0a16}/src/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-8.0.0a14 → arize-8.0.0a16}/src/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-8.0.0a14 → arize-8.0.0a16}/src/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,