arize 8.0.0a12__tar.gz → 8.0.0a14__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize
-Version: 8.0.0a12
+Version: 8.0.0a14
 Summary: A helper library to interact with Arize AI APIs
 Project-URL: Homepage, https://arize.com
 Project-URL: Documentation, https://docs.arize.com/arize
@@ -34,6 +34,8 @@ Requires-Dist: pillow<11,>=8.4.0; extra == 'auto-embeddings'
 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: pydantic; extra == 'datasets-experiments'
 Provides-Extra: dev
 Requires-Dist: pytest==8.4.2; extra == 'dev'
 Requires-Dist: ruff==0.13.2; extra == 'dev'
@@ -92,6 +94,11 @@ Description-Content-Type: text/markdown
     - [Log a batch of ML Data for a Object Detection use-case](#log-a-batch-of-ml-data-for-a-object-detection-use-case)
     - [Exporting ML Data](#exporting-ml-data)
   - [Generate embeddings for your data](#generate-embeddings-for-your-data)
+  - [Operations on Datasets](#operations-on-datasets)
+    - [List Datasets](#list-datasets)
+    - [Create a Dataset](#create-a-dataset)
+    - [Get Dataset by ID](#get-dataset-by-id)
+    - [Delete a Dataset](#delete-a-dataset)
 - [Configure Logging](#configure-logging)
   - [In Code](#in-code)
   - [Via Environment Variables](#via-environment-variables)
@@ -368,6 +375,96 @@ generator = EmbeddingGenerator.from_use_case(
 df["text_vector"] = generator.generate_embeddings(text_col=df["text"])
 ```
 
+## Operations on Datasets
+
+### List Datasets
+
+You can list all datasets that the user has access to using `client.datasets.list()`. You can use the `limit` parameter to specify the maximum number of datasets desired in the response and you can specify the `space_id` to target the list operation to a particular space.
+
+```python
+resp = client.datasets.list(
+    limit=... # Optional
+    space_id=... # Optional
+)
+```
+
+The response is an object of type `DatasetsList200Response`, and you can access the list of datasets via its `datasets` attribute. In addition, you can transform the response object to a dictionary, to JSON format, or a pandas dataframe.
+
+```python
+# Get the list of datasets from the response
+dataset_list = resp.datasets 
+# Get the response as a dictionary
+resp_dict = resp.to_dict()
+# Get the response in JSON format
+resp_dict = resp.to_json()
+# Get the response as a pandas dataframe
+resp_dict = resp.to_df()
+```
+
+### Create a Dataset
+
+You can create a dataset using `client.datasets.create()`. You must pass examples, we currently don't support creating an empty dataset, for instance, these are 2 rows of examples, as a list of dictionaries. You can also pass a pandas dataframe for the examples.
+
+```python
+examples = [
+    {
+        "eval.Correctness Basic.explanation": "The query indicates that the user is having trouble accessing their account on their laptop, while access on their phone is still working. This suggests a potential issue with the login process on the laptop, which aligns with the 'Login Issues' queue. The mention of a possible change in the account could relate to login credentials or settings affecting the laptop specifically, but it still falls under the broader category of login issues.",
+        "eval.Correctness Basic.label": "correct",
+        "eval.Correctness Basic.score": 1,
+        "llm output": "Login Issues",
+        "query": "I can't get in on my laptop anymore, but my phone still works fine — could this be because I changed something in my account?"
+    },
+    {
+        "eval.Correctness Basic.explanation": "The query is about a user who signed up but is unable to log in because the system says no account is found. This issue is related to the login process, as the user is trying to access their account and is facing a problem with the login system recognizing their account. Therefore, assigning this query to the 'Login Issues' queue is appropriate.",
+        "eval.Correctness Basic.label": "correct",
+        "eval.Correctness Basic.score": 1,
+        "llm output": "Login Issues",
+        "query": "Signed up ages ago but never got around to logging in — now it says no account found. Do I start over?"
+    }
+]
+```
+
+If the number of examples (rows in dataframe, items in list) is too large, the client SDK will try to send the data via Arrow Flight via gRPC for better performance. If you want to force the data transfer to HTTP you can use the `force_http` flag. The response is a `Dataset` object.
+
+```python
+created_dataset = client.datasets.create(
+    space_i="<target-space-id>",
+    name="<your-dataset-name>", # Name must be unique within a space
+    examples=..., # List of dictionaries or pandas dataframe
+)
+```
+
+The `Dataset` object also counts with convenience method similar to `List***` objects:
+
+```python
+# Get the response as a dictionary
+dataset_dict = create_dataset.to_dict()
+# Get the response in JSON format
+dataset_dict = create_dataset.to_json()
+```
+
+
+### Get Dataset by ID
+
+To get a dataset by its ID use `client.datasets.get()`, you can optionally also pass the version ID of a particular version of interest of the dataset. The returned type is `Dataset`.
+
+```python
+dataset = client.datasets.get(
+    dataset_id=... # The unique identifier of the dataset
+    dataset_version_id=... # The unique identifier of the dataset version
+)
+```
+
+### Delete a Dataset
+
+To delete a dataset by its ID use `client.datasets.delete()`. The call returns `None` if successful deletion took place, error otherwise.
+
+```python
+client.datasets.delete(
+    dataset_id=... # The unique identifier of the dataset
+)
+```
+
 # Configure Logging
 
 ## In Code

{arize-8.0.0a12 → arize-8.0.0a14}/README.md

@@ -32,6 +32,11 @@
     - [Log a batch of ML Data for a Object Detection use-case](#log-a-batch-of-ml-data-for-a-object-detection-use-case)
     - [Exporting ML Data](#exporting-ml-data)
   - [Generate embeddings for your data](#generate-embeddings-for-your-data)
+  - [Operations on Datasets](#operations-on-datasets)
+    - [List Datasets](#list-datasets)
+    - [Create a Dataset](#create-a-dataset)
+    - [Get Dataset by ID](#get-dataset-by-id)
+    - [Delete a Dataset](#delete-a-dataset)
 - [Configure Logging](#configure-logging)
   - [In Code](#in-code)
   - [Via Environment Variables](#via-environment-variables)
@@ -308,6 +313,96 @@ generator = EmbeddingGenerator.from_use_case(
 df["text_vector"] = generator.generate_embeddings(text_col=df["text"])
 ```
 
+## Operations on Datasets
+
+### List Datasets
+
+You can list all datasets that the user has access to using `client.datasets.list()`. You can use the `limit` parameter to specify the maximum number of datasets desired in the response and you can specify the `space_id` to target the list operation to a particular space.
+
+```python
+resp = client.datasets.list(
+    limit=... # Optional
+    space_id=... # Optional
+)
+```
+
+The response is an object of type `DatasetsList200Response`, and you can access the list of datasets via its `datasets` attribute. In addition, you can transform the response object to a dictionary, to JSON format, or a pandas dataframe.
+
+```python
+# Get the list of datasets from the response
+dataset_list = resp.datasets 
+# Get the response as a dictionary
+resp_dict = resp.to_dict()
+# Get the response in JSON format
+resp_dict = resp.to_json()
+# Get the response as a pandas dataframe
+resp_dict = resp.to_df()
+```
+
+### Create a Dataset
+
+You can create a dataset using `client.datasets.create()`. You must pass examples, we currently don't support creating an empty dataset, for instance, these are 2 rows of examples, as a list of dictionaries. You can also pass a pandas dataframe for the examples.
+
+```python
+examples = [
+    {
+        "eval.Correctness Basic.explanation": "The query indicates that the user is having trouble accessing their account on their laptop, while access on their phone is still working. This suggests a potential issue with the login process on the laptop, which aligns with the 'Login Issues' queue. The mention of a possible change in the account could relate to login credentials or settings affecting the laptop specifically, but it still falls under the broader category of login issues.",
+        "eval.Correctness Basic.label": "correct",
+        "eval.Correctness Basic.score": 1,
+        "llm output": "Login Issues",
+        "query": "I can't get in on my laptop anymore, but my phone still works fine — could this be because I changed something in my account?"
+    },
+    {
+        "eval.Correctness Basic.explanation": "The query is about a user who signed up but is unable to log in because the system says no account is found. This issue is related to the login process, as the user is trying to access their account and is facing a problem with the login system recognizing their account. Therefore, assigning this query to the 'Login Issues' queue is appropriate.",
+        "eval.Correctness Basic.label": "correct",
+        "eval.Correctness Basic.score": 1,
+        "llm output": "Login Issues",
+        "query": "Signed up ages ago but never got around to logging in — now it says no account found. Do I start over?"
+    }
+]
+```
+
+If the number of examples (rows in dataframe, items in list) is too large, the client SDK will try to send the data via Arrow Flight via gRPC for better performance. If you want to force the data transfer to HTTP you can use the `force_http` flag. The response is a `Dataset` object.
+
+```python
+created_dataset = client.datasets.create(
+    space_i="<target-space-id>",
+    name="<your-dataset-name>", # Name must be unique within a space
+    examples=..., # List of dictionaries or pandas dataframe
+)
+```
+
+The `Dataset` object also counts with convenience method similar to `List***` objects:
+
+```python
+# Get the response as a dictionary
+dataset_dict = create_dataset.to_dict()
+# Get the response in JSON format
+dataset_dict = create_dataset.to_json()
+```
+
+
+### Get Dataset by ID
+
+To get a dataset by its ID use `client.datasets.get()`, you can optionally also pass the version ID of a particular version of interest of the dataset. The returned type is `Dataset`.
+
+```python
+dataset = client.datasets.get(
+    dataset_id=... # The unique identifier of the dataset
+    dataset_version_id=... # The unique identifier of the dataset version
+)
+```
+
+### Delete a Dataset
+
+To delete a dataset by its ID use `client.datasets.delete()`. The call returns `None` if successful deletion took place, error otherwise.
+
+```python
+client.datasets.delete(
+    dataset_id=... # The unique identifier of the dataset
+)
+```
+
 # Configure Logging
 
 ## In Code

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

@@ -70,9 +70,9 @@ ml-batch = [
     "requests>=2.0.0, <3", # For posting pyarrow files
     "tqdm", # For export progress bars
 ]
-# datasets-experiments = [
-#     "pydantic",
-# ]
+datasets-experiments = [
+    "pydantic",
+]
 mimic-explainer = [
     "interpret-community[mimic]>=0.22.0,<1",
 ]

{arize-8.0.0a12 → arize-8.0.0a14}/src/arize/_exporter/client.py

@@ -20,7 +20,6 @@ from arize._generated.protocol.flight import export_pb2
 from arize.logging import CtxAdapter
 from arize.types import Environments, SimilaritySearchParams
 from arize.utils.dataframe import reset_dataframe_index
-from arize.utils.proto import get_pb_flight_doput_request
 
 logger = logging.getLogger(__name__)
 
@@ -131,7 +130,7 @@ class ArizeExportClient:
         reset_dataframe_index(df)
         return df
 
-    def export_model_to_parquet(
+    def export_to_parquet(
         self,
         path: str,
         space_id: str,
@@ -285,7 +284,7 @@ class ArizeExportClient:
             end_time=Timestamp(seconds=int(end_time.timestamp())),
             filter_expression=where,
             similarity_search_params=(
-                get_pb_flight_doput_request(similarity_search_params)
+                _get_pb_similarity_search_params(similarity_search_params)
                 if similarity_search_params
                 else None
             ),
@@ -326,3 +325,19 @@ class ArizeExportClient:
             colour="#008000",
             unit=" row",
         )
+
+
+def _get_pb_similarity_search_params(
+    similarity_params: SimilaritySearchParams,
+) -> export_pb2.SimilaritySearchParams:
+    proto_params = export_pb2.SimilaritySearchParams()
+    proto_params.search_column_name = similarity_params.search_column_name
+    proto_params.threshold = similarity_params.threshold
+    for ref in similarity_params.references:
+        new_ref = proto_params.references.add()
+        new_ref.prediction_id = ref.prediction_id
+        new_ref.reference_column_name = ref.reference_column_name
+        if ref.prediction_timestamp:
+            new_ref.prediction_timestamp.FromDatetime(ref.prediction_timestamp)
+
+    return proto_params

{arize-8.0.0a12 → arize-8.0.0a14}/src/arize/_flight/client.py

@@ -10,6 +10,7 @@ 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 (
     WriteSpanAnnotationResponse,
     WriteSpanAttributesMetadataResponse,
@@ -17,7 +18,6 @@ from arize._generated.protocol.flight.ingest_pb2 import (
 )
 from arize.config import get_python_version
 from arize.logging import log_a_list
-from arize.utils.proto import get_pb_flight_doput_request, get_pb_schema_tracing
 from arize.version import __version__
 
 if TYPE_CHECKING:
@@ -137,7 +137,7 @@ class ArizeFlightClient:
             pa_table.schema, {"arize-schema": base64_schema}
         )
 
-        doput_request = get_pb_flight_doput_request(
+        doput_request = _get_pb_flight_doput_request(
             space_id=space_id,
             model_id=project_name,
             request_type=request_type,
@@ -179,7 +179,49 @@ class ArizeFlightClient:
                 return res
         except Exception as e:
             logger.exception(f"Error logging arrow table to Arize: {e}")
-            raise
+                    raise RuntimeError(f"Error logging arrow table to Arize: {e}") from e
+
+    # ---------- dataset methods ----------
+
+    def create_dataset(
+        self,
+        space_id: str,
+        dataset_name: str,
+        pa_table: pa.Table,
+    ) -> str:
+        doput_request = flight_ing_pb2.DoPutRequest(
+            create_dataset=flight_ing_pb2.CreateDatasetRequest(
+                space_id=space_id,
+                dataset_name=dataset_name,
+                dataset_type=flight_ing_pb2.GENERATIVE,
+            )
+        )
+        descriptor = flight.FlightDescriptor.for_command(
+            json_format.MessageToJson(doput_request).encode("utf-8")
+        )
+        try:
+            flight_writer, flight_metadata_reader = self.do_put(
+                descriptor, pa_table.schema, options=self.call_options
+            )
+            with flight_writer:
+                # write table as stream to flight server
+                flight_writer.write_table(pa_table)
+                # indicate that client has flushed all contents to stream
+                flight_writer.done_writing()
+                # read response from flight server
+                flight_response = flight_metadata_reader.read()
+                if flight_response is None:
+                    return None
+
+                res = None
+                res = flight_ing_pb2.CreateDatasetResponse()
+                res.ParseFromString(flight_response.to_pybytes())
+                if res:
+                    return str(res.dataset_id)
+                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
 
 
 def append_to_pyarrow_metadata(
@@ -201,3 +243,37 @@ def append_to_pyarrow_metadata(
     updated_metadata = metadata.copy()
     updated_metadata.update(new_metadata)
     return pa_schema.with_metadata(updated_metadata)
+
+
+def _get_pb_flight_doput_request(
+    space_id,
+    model_id: str,
+    request_type: FlightRequestType,
+) -> 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
+                )
+            )
+        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 _:
+            raise ValueError(f"Unsupported request_type: {request_type}")

{arize-8.0.0a12 → arize-8.0.0a14}/src/arize/_generated/api_client/__init__.py

@@ -31,7 +31,6 @@ __all__ = [
     "ApiException",
     "Dataset",
     "DatasetVersion",
-    "DatasetsCreate201Response",
     "DatasetsCreateRequest",
     "DatasetsList200Response",
     "DatasetsListExamples200Response",
@@ -40,42 +39,7 @@ __all__ = [
     "ExperimentsList200Response",
 ]
 
-if __import__("typing").TYPE_CHECKING:
-    # 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
-    
-    # import ApiClient
-    from arize._generated.api_client.api_response import ApiResponse as ApiResponse
-    from arize._generated.api_client.api_client import ApiClient as ApiClient
-    from arize._generated.api_client.configuration import Configuration as Configuration
-    from arize._generated.api_client.exceptions import OpenApiException as OpenApiException
-    from arize._generated.api_client.exceptions import ApiTypeError as ApiTypeError
-    from arize._generated.api_client.exceptions import ApiValueError as ApiValueError
-    from arize._generated.api_client.exceptions import ApiKeyError as ApiKeyError
-    from arize._generated.api_client.exceptions import ApiAttributeError as ApiAttributeError
-    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_version import DatasetVersion as DatasetVersion
-    from arize._generated.api_client.models.datasets_create201_response import DatasetsCreate201Response as DatasetsCreate201Response
-    from arize._generated.api_client.models.datasets_create_request import DatasetsCreateRequest as DatasetsCreateRequest
-    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.experiments_list200_response import ExperimentsList200Response as ExperimentsList200Response
-    
-else:
-    from lazy_imports import LazyModule, as_package, load
-
-    load(
-        LazyModule(
-            *as_package(__file__),
-            ("__version__", __version__),
-            ("__all__", __all__),
-            """# import apis into sdk package
+# 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
 
@@ -93,7 +57,6 @@ 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_version import DatasetVersion as DatasetVersion
-from arize._generated.api_client.models.datasets_create201_response import DatasetsCreate201Response as DatasetsCreate201Response
 from arize._generated.api_client.models.datasets_create_request import DatasetsCreateRequest as DatasetsCreateRequest
 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
@@ -101,8 +64,3 @@ 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_list200_response import ExperimentsList200Response as ExperimentsList200Response
 
-""",
-            name=__name__,
-            doc=__doc__,
-        )
-    )

arize-8.0.0a14/src/arize/_generated/api_client/api/__init__.py

@@ -0,0 +1,6 @@
+# flake8: noqa
+
+# 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
+

{arize-8.0.0a12 → arize-8.0.0a14}/src/arize/_generated/api_client/api/datasets_api.py

@@ -20,7 +20,6 @@ from pydantic import Field, StrictInt, StrictStr
 from typing import Optional
 from typing_extensions import Annotated
 from arize._generated.api_client.models.dataset import Dataset
-from arize._generated.api_client.models.datasets_create201_response import DatasetsCreate201Response
 from arize._generated.api_client.models.datasets_create_request import DatasetsCreateRequest
 from arize._generated.api_client.models.datasets_list200_response import DatasetsList200Response
 from arize._generated.api_client.models.datasets_list_examples200_response import DatasetsListExamples200Response
@@ -59,7 +58,7 @@ class DatasetsApi:
         _content_type: Optional[StrictStr] = None,
         _headers: Optional[Dict[StrictStr, Any]] = None,
         _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
-    ) -> DatasetsCreate201Response:
+    ) -> Dataset:
         """Create a new dataset with JSON examples
 
 
@@ -96,7 +95,7 @@ class DatasetsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '201': "DatasetsCreate201Response",
+            '201': "Dataset",
             '400': "Error",
             '401': "Error",
             '403': "Error",
@@ -129,7 +128,7 @@ class DatasetsApi:
         _content_type: Optional[StrictStr] = None,
         _headers: Optional[Dict[StrictStr, Any]] = None,
         _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
-    ) -> ApiResponse[DatasetsCreate201Response]:
+    ) -> ApiResponse[Dataset]:
         """Create a new dataset with JSON examples
 
 
@@ -166,7 +165,7 @@ class DatasetsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '201': "DatasetsCreate201Response",
+            '201': "Dataset",
             '400': "Error",
             '401': "Error",
             '403': "Error",
@@ -236,7 +235,7 @@ class DatasetsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '201': "DatasetsCreate201Response",
+            '201': "Dataset",
             '400': "Error",
             '401': "Error",
             '403': "Error",
@@ -342,7 +341,7 @@ class DatasetsApi:
         _content_type: Optional[StrictStr] = None,
         _headers: Optional[Dict[StrictStr, Any]] = None,
         _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
-    ) -> DatasetsCreate201Response:
+    ) -> None:
         """Delete a dataset by ID
 
 
@@ -379,7 +378,7 @@ class DatasetsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '200': "DatasetsCreate201Response",
+            '204': None,
             '400': "Error",
             '401': "Error",
             '403': "Error",
@@ -414,7 +413,7 @@ class DatasetsApi:
         _content_type: Optional[StrictStr] = None,
         _headers: Optional[Dict[StrictStr, Any]] = None,
         _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
-    ) -> ApiResponse[DatasetsCreate201Response]:
+    ) -> ApiResponse[None]:
         """Delete a dataset by ID
 
 
@@ -451,7 +450,7 @@ class DatasetsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '200': "DatasetsCreate201Response",
+            '204': None,
             '400': "Error",
             '401': "Error",
             '403': "Error",
@@ -523,7 +522,7 @@ class DatasetsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '200': "DatasetsCreate201Response",
+            '204': None,
             '400': "Error",
             '401': "Error",
             '403': "Error",