arize 8.0.0a11__py3-none-any.whl → 8.0.0a13__py3-none-any.whl

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/_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/_generated/api_client/api/__init__.py

@@ -1,22 +1,6 @@
 # flake8: noqa
 
-if __import__("typing").TYPE_CHECKING:
-    # 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
-    
-else:
-    from lazy_imports import LazyModule, as_package, load
-
-    load(
-        LazyModule(
-            *as_package(__file__),
-            """# import apis into api package
+# 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
 
-""",
-            name=__name__,
-            doc=__doc__,
-        )
-    )

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",

arize/_generated/api_client/api/experiments_api.py

@@ -39,6 +39,282 @@ class ExperimentsApi:
         self.api_client = api_client
 
 
+    @validate_call
+    def experiments_delete(
+        self,
+        experiment_id: Annotated[StrictStr, Field(description="The unique identifier of the experiment")],
+        _request_timeout: Union[
+            None,
+            Annotated[StrictFloat, Field(gt=0)],
+            Tuple[
+                Annotated[StrictFloat, Field(gt=0)],
+                Annotated[StrictFloat, Field(gt=0)]
+            ]
+        ] = None,
+        _request_auth: Optional[Dict[StrictStr, Any]] = None,
+        _content_type: Optional[StrictStr] = None,
+        _headers: Optional[Dict[StrictStr, Any]] = None,
+        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
+    ) -> None:
+        """Delete an experiment by ID
+
+
+        :param experiment_id: The unique identifier of the experiment (required)
+        :type experiment_id: str
+        :param _request_timeout: timeout setting for this request. If one
+                                 number provided, it will be total request
+                                 timeout. It can also be a pair (tuple) of
+                                 (connection, read) timeouts.
+        :type _request_timeout: int, tuple(int, int), optional
+        :param _request_auth: set to override the auth_settings for an a single
+                              request; this effectively ignores the
+                              authentication in the spec for a single request.
+        :type _request_auth: dict, optional
+        :param _content_type: force content-type for the request.
+        :type _content_type: str, Optional
+        :param _headers: set to override the headers for a single
+                         request; this effectively ignores the headers
+                         in the spec for a single request.
+        :type _headers: dict, optional
+        :param _host_index: set to override the host_index for a single
+                            request; this effectively ignores the host_index
+                            in the spec for a single request.
+        :type _host_index: int, optional
+        :return: Returns the result object.
+        """ # noqa: E501
+
+        _param = self._experiments_delete_serialize(
+            experiment_id=experiment_id,
+            _request_auth=_request_auth,
+            _content_type=_content_type,
+            _headers=_headers,
+            _host_index=_host_index
+        )
+
+        _response_types_map: Dict[str, Optional[str]] = {
+            '204': None,
+            '400': "Error",
+            '401': "Error",
+            '403': "Error",
+            '404': "Error",
+            '422': "Error",
+            '429': "Error",
+        }
+        response_data = self.api_client.call_api(
+            *_param,
+            _request_timeout=_request_timeout
+        )
+        response_data.read()
+        return self.api_client.response_deserialize(
+            response_data=response_data,
+            response_types_map=_response_types_map,
+        ).data
+
+
+    @validate_call
+    def experiments_delete_with_http_info(
+        self,
+        experiment_id: Annotated[StrictStr, Field(description="The unique identifier of the experiment")],
+        _request_timeout: Union[
+            None,
+            Annotated[StrictFloat, Field(gt=0)],
+            Tuple[
+                Annotated[StrictFloat, Field(gt=0)],
+                Annotated[StrictFloat, Field(gt=0)]
+            ]
+        ] = None,
+        _request_auth: Optional[Dict[StrictStr, Any]] = None,
+        _content_type: Optional[StrictStr] = None,
+        _headers: Optional[Dict[StrictStr, Any]] = None,
+        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
+    ) -> ApiResponse[None]:
+        """Delete an experiment by ID
+
+
+        :param experiment_id: The unique identifier of the experiment (required)
+        :type experiment_id: str
+        :param _request_timeout: timeout setting for this request. If one
+                                 number provided, it will be total request
+                                 timeout. It can also be a pair (tuple) of
+                                 (connection, read) timeouts.
+        :type _request_timeout: int, tuple(int, int), optional
+        :param _request_auth: set to override the auth_settings for an a single
+                              request; this effectively ignores the
+                              authentication in the spec for a single request.
+        :type _request_auth: dict, optional
+        :param _content_type: force content-type for the request.
+        :type _content_type: str, Optional
+        :param _headers: set to override the headers for a single
+                         request; this effectively ignores the headers
+                         in the spec for a single request.
+        :type _headers: dict, optional
+        :param _host_index: set to override the host_index for a single
+                            request; this effectively ignores the host_index
+                            in the spec for a single request.
+        :type _host_index: int, optional
+        :return: Returns the result object.
+        """ # noqa: E501
+
+        _param = self._experiments_delete_serialize(
+            experiment_id=experiment_id,
+            _request_auth=_request_auth,
+            _content_type=_content_type,
+            _headers=_headers,
+            _host_index=_host_index
+        )
+
+        _response_types_map: Dict[str, Optional[str]] = {
+            '204': None,
+            '400': "Error",
+            '401': "Error",
+            '403': "Error",
+            '404': "Error",
+            '422': "Error",
+            '429': "Error",
+        }
+        response_data = self.api_client.call_api(
+            *_param,
+            _request_timeout=_request_timeout
+        )
+        response_data.read()
+        return self.api_client.response_deserialize(
+            response_data=response_data,
+            response_types_map=_response_types_map,
+        )
+
+
+    @validate_call
+    def experiments_delete_without_preload_content(
+        self,
+        experiment_id: Annotated[StrictStr, Field(description="The unique identifier of the experiment")],
+        _request_timeout: Union[
+            None,
+            Annotated[StrictFloat, Field(gt=0)],
+            Tuple[
+                Annotated[StrictFloat, Field(gt=0)],
+                Annotated[StrictFloat, Field(gt=0)]
+            ]
+        ] = None,
+        _request_auth: Optional[Dict[StrictStr, Any]] = None,
+        _content_type: Optional[StrictStr] = None,
+        _headers: Optional[Dict[StrictStr, Any]] = None,
+        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
+    ) -> RESTResponseType:
+        """Delete an experiment by ID
+
+
+        :param experiment_id: The unique identifier of the experiment (required)
+        :type experiment_id: str
+        :param _request_timeout: timeout setting for this request. If one
+                                 number provided, it will be total request
+                                 timeout. It can also be a pair (tuple) of
+                                 (connection, read) timeouts.
+        :type _request_timeout: int, tuple(int, int), optional
+        :param _request_auth: set to override the auth_settings for an a single
+                              request; this effectively ignores the
+                              authentication in the spec for a single request.
+        :type _request_auth: dict, optional
+        :param _content_type: force content-type for the request.
+        :type _content_type: str, Optional
+        :param _headers: set to override the headers for a single
+                         request; this effectively ignores the headers
+                         in the spec for a single request.
+        :type _headers: dict, optional
+        :param _host_index: set to override the host_index for a single
+                            request; this effectively ignores the host_index
+                            in the spec for a single request.
+        :type _host_index: int, optional
+        :return: Returns the result object.
+        """ # noqa: E501
+
+        _param = self._experiments_delete_serialize(
+            experiment_id=experiment_id,
+            _request_auth=_request_auth,
+            _content_type=_content_type,
+            _headers=_headers,
+            _host_index=_host_index
+        )
+
+        _response_types_map: Dict[str, Optional[str]] = {
+            '204': None,
+            '400': "Error",
+            '401': "Error",
+            '403': "Error",
+            '404': "Error",
+            '422': "Error",
+            '429': "Error",
+        }
+        response_data = self.api_client.call_api(
+            *_param,
+            _request_timeout=_request_timeout
+        )
+        return response_data.response
+
+
+    def _experiments_delete_serialize(
+        self,
+        experiment_id,
+        _request_auth,
+        _content_type,
+        _headers,
+        _host_index,
+    ) -> RequestSerialized:
+
+        _host = None
+
+        _collection_formats: Dict[str, str] = {
+        }
+
+        _path_params: Dict[str, str] = {}
+        _query_params: List[Tuple[str, str]] = []
+        _header_params: Dict[str, Optional[str]] = _headers or {}
+        _form_params: List[Tuple[str, str]] = []
+        _files: Dict[
+            str, Union[str, bytes, List[str], List[bytes], List[Tuple[str, bytes]]]
+        ] = {}
+        _body_params: Optional[bytes] = None
+
+        # process the path parameters
+        if experiment_id is not None:
+            _path_params['experimentId'] = experiment_id
+        # process the query parameters
+        # process the header parameters
+        # process the form parameters
+        # process the body parameter
+
+
+        # set the HTTP header `Accept`
+        if 'Accept' not in _header_params:
+            _header_params['Accept'] = self.api_client.select_header_accept(
+                [
+                    'application/json'
+                ]
+            )
+
+
+        # authentication setting
+        _auth_settings: List[str] = [
+            'ApiKeyAuth'
+        ]
+
+        return self.api_client.param_serialize(
+            method='DELETE',
+            resource_path='/api/v1/experiments/{experimentId}',
+            path_params=_path_params,
+            query_params=_query_params,
+            header_params=_header_params,
+            body=_body_params,
+            post_params=_form_params,
+            files=_files,
+            auth_settings=_auth_settings,
+            collection_formats=_collection_formats,
+            _host=_host,
+            _request_auth=_request_auth
+        )
+
+
+
+
     @validate_call
     def experiments_list(
         self,

arize/_generated/api_client/api_client.py

@@ -460,13 +460,13 @@ class ApiClient:
 
         if klass in self.PRIMITIVE_TYPES:
             return self.__deserialize_primitive(data, klass)
-        elif klass == object:
+        elif klass is object:
             return self.__deserialize_object(data)
-        elif klass == datetime.date:
+        elif klass is datetime.date:
             return self.__deserialize_date(data)
-        elif klass == datetime.datetime:
+        elif klass is datetime.datetime:
             return self.__deserialize_datetime(data)
-        elif klass == decimal.Decimal:
+        elif klass is decimal.Decimal:
             return decimal.Decimal(data)
         elif issubclass(klass, Enum):
             return self.__deserialize_enum(data, klass)

arize/_generated/api_client/models/__init__.py

@@ -12,29 +12,9 @@
     Do not edit the class manually.
 """  # noqa: E501
 
-
-if __import__("typing").TYPE_CHECKING:
-    # import models into model package
-    from arize._generated.api_client.models.dataset import Dataset
-    from arize._generated.api_client.models.dataset_version import DatasetVersion
-    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
-    from arize._generated.api_client.models.error import Error
-    from arize._generated.api_client.models.experiment import Experiment
-    from arize._generated.api_client.models.experiments_list200_response import ExperimentsList200Response
-    
-else:
-    from lazy_imports import LazyModule, as_package, load
-
-    load(
-        LazyModule(
-            *as_package(__file__),
-            """# import models into model package
+# import models into model package
 from arize._generated.api_client.models.dataset import Dataset
 from arize._generated.api_client.models.dataset_version import DatasetVersion
-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
@@ -42,8 +22,3 @@ from arize._generated.api_client.models.error import Error
 from arize._generated.api_client.models.experiment import Experiment
 from arize._generated.api_client.models.experiments_list200_response import ExperimentsList200Response
 
-""",
-            name=__name__,
-            doc=__doc__,
-        )
-    )