opik 1.8.39__py3-none-any.whl → 1.9.71__py3-none-any.whl

opik/rest_api/optimizations/raw_client.py

@@ -10,11 +10,15 @@ from ..core.http_response import AsyncHttpResponse, HttpResponse
 from ..core.jsonable_encoder import jsonable_encoder
 from ..core.pydantic_utilities import parse_obj_as
 from ..core.request_options import RequestOptions
+from ..core.serialization import convert_and_respect_annotation_metadata
 from ..errors.bad_request_error import BadRequestError
 from ..errors.not_found_error import NotFoundError
-from ..types.json_node_write import JsonNodeWrite
+from ..errors.not_implemented_error import NotImplementedError
+from ..types.json_list_string_write import JsonListStringWrite
 from ..types.optimization_page_public import OptimizationPagePublic
 from ..types.optimization_public import OptimizationPublic
+from ..types.optimization_studio_config_write import OptimizationStudioConfigWrite
+from ..types.optimization_studio_log import OptimizationStudioLog
 from ..types.optimization_write_status import OptimizationWriteStatus
 from .types.optimization_update_status import OptimizationUpdateStatus
 
@@ -26,6 +30,47 @@ class RawOptimizationsClient:
     def __init__(self, *, client_wrapper: SyncClientWrapper):
         self._client_wrapper = client_wrapper
 
+    def cancel_studio_optimizations(
+        self, id: str, *, request_options: typing.Optional[RequestOptions] = None
+    ) -> HttpResponse[None]:
+        """
+        Cancel Studio optimizations by id
+
+        Parameters
+        ----------
+        id : str
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        HttpResponse[None]
+        """
+        _response = self._client_wrapper.httpx_client.request(
+            f"v1/private/optimizations/studio/{jsonable_encoder(id)}/cancel",
+            method="GET",
+            request_options=request_options,
+        )
+        try:
+            if 200 <= _response.status_code < 300:
+                return HttpResponse(response=_response, data=None)
+            if _response.status_code == 501:
+                raise NotImplementedError(
+                    headers=dict(_response.headers),
+                    body=typing.cast(
+                        typing.Optional[typing.Any],
+                        parse_obj_as(
+                            type_=typing.Optional[typing.Any],  # type: ignore
+                            object_=_response.json(),
+                        ),
+                    ),
+                )
+            _response_json = _response.json()
+        except JSONDecodeError:
+            raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
     def find_optimizations(
         self,
         *,
@@ -34,6 +79,7 @@ class RawOptimizationsClient:
         dataset_id: typing.Optional[str] = None,
         name: typing.Optional[str] = None,
         dataset_deleted: typing.Optional[bool] = None,
+        filters: typing.Optional[str] = None,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> HttpResponse[OptimizationPagePublic]:
         """
@@ -51,6 +97,8 @@ class RawOptimizationsClient:
 
         dataset_deleted : typing.Optional[bool]
 
+        filters : typing.Optional[str]
+
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
@@ -68,6 +116,7 @@ class RawOptimizationsClient:
                 "dataset_id": dataset_id,
                 "name": name,
                 "dataset_deleted": dataset_deleted,
+                "filters": filters,
             },
             request_options=request_options,
         )
@@ -105,7 +154,8 @@ class RawOptimizationsClient:
         status: OptimizationWriteStatus,
         id: typing.Optional[str] = OMIT,
         name: typing.Optional[str] = OMIT,
-        metadata: typing.Optional[JsonNodeWrite] = OMIT,
+        metadata: typing.Optional[JsonListStringWrite] = OMIT,
+        studio_config: typing.Optional[OptimizationStudioConfigWrite] = OMIT,
         last_updated_at: typing.Optional[dt.datetime] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> HttpResponse[None]:
@@ -124,7 +174,9 @@ class RawOptimizationsClient:
 
         name : typing.Optional[str]
 
-        metadata : typing.Optional[JsonNodeWrite]
+        metadata : typing.Optional[JsonListStringWrite]
+
+        studio_config : typing.Optional[OptimizationStudioConfigWrite]
 
         last_updated_at : typing.Optional[dt.datetime]
 
@@ -144,7 +196,12 @@ class RawOptimizationsClient:
                 "dataset_name": dataset_name,
                 "objective_name": objective_name,
                 "status": status,
-                "metadata": metadata,
+                "metadata": convert_and_respect_annotation_metadata(
+                    object_=metadata, annotation=JsonListStringWrite, direction="write"
+                ),
+                "studio_config": convert_and_respect_annotation_metadata(
+                    object_=studio_config, annotation=OptimizationStudioConfigWrite, direction="write"
+                ),
                 "last_updated_at": last_updated_at,
             },
             headers={
@@ -169,7 +226,8 @@ class RawOptimizationsClient:
         status: OptimizationWriteStatus,
         id: typing.Optional[str] = OMIT,
         name: typing.Optional[str] = OMIT,
-        metadata: typing.Optional[JsonNodeWrite] = OMIT,
+        metadata: typing.Optional[JsonListStringWrite] = OMIT,
+        studio_config: typing.Optional[OptimizationStudioConfigWrite] = OMIT,
         last_updated_at: typing.Optional[dt.datetime] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> HttpResponse[None]:
@@ -188,7 +246,9 @@ class RawOptimizationsClient:
 
         name : typing.Optional[str]
 
-        metadata : typing.Optional[JsonNodeWrite]
+        metadata : typing.Optional[JsonListStringWrite]
+
+        studio_config : typing.Optional[OptimizationStudioConfigWrite]
 
         last_updated_at : typing.Optional[dt.datetime]
 
@@ -208,7 +268,12 @@ class RawOptimizationsClient:
                 "dataset_name": dataset_name,
                 "objective_name": objective_name,
                 "status": status,
-                "metadata": metadata,
+                "metadata": convert_and_respect_annotation_metadata(
+                    object_=metadata, annotation=JsonListStringWrite, direction="write"
+                ),
+                "studio_config": convert_and_respect_annotation_metadata(
+                    object_=studio_config, annotation=OptimizationStudioConfigWrite, direction="write"
+                ),
                 "last_updated_at": last_updated_at,
             },
             headers={
@@ -358,11 +423,101 @@ class RawOptimizationsClient:
             raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
         raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
 
+    def get_studio_optimization_logs(
+        self, id: str, *, request_options: typing.Optional[RequestOptions] = None
+    ) -> HttpResponse[OptimizationStudioLog]:
+        """
+        Get presigned S3 URL for downloading optimization logs
+
+        Parameters
+        ----------
+        id : str
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        HttpResponse[OptimizationStudioLog]
+            Logs response
+        """
+        _response = self._client_wrapper.httpx_client.request(
+            f"v1/private/optimizations/studio/{jsonable_encoder(id)}/logs",
+            method="GET",
+            request_options=request_options,
+        )
+        try:
+            if 200 <= _response.status_code < 300:
+                _data = typing.cast(
+                    OptimizationStudioLog,
+                    parse_obj_as(
+                        type_=OptimizationStudioLog,  # type: ignore
+                        object_=_response.json(),
+                    ),
+                )
+                return HttpResponse(response=_response, data=_data)
+            if _response.status_code == 404:
+                raise NotFoundError(
+                    headers=dict(_response.headers),
+                    body=typing.cast(
+                        typing.Optional[typing.Any],
+                        parse_obj_as(
+                            type_=typing.Optional[typing.Any],  # type: ignore
+                            object_=_response.json(),
+                        ),
+                    ),
+                )
+            _response_json = _response.json()
+        except JSONDecodeError:
+            raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
 
 class AsyncRawOptimizationsClient:
     def __init__(self, *, client_wrapper: AsyncClientWrapper):
         self._client_wrapper = client_wrapper
 
+    async def cancel_studio_optimizations(
+        self, id: str, *, request_options: typing.Optional[RequestOptions] = None
+    ) -> AsyncHttpResponse[None]:
+        """
+        Cancel Studio optimizations by id
+
+        Parameters
+        ----------
+        id : str
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        AsyncHttpResponse[None]
+        """
+        _response = await self._client_wrapper.httpx_client.request(
+            f"v1/private/optimizations/studio/{jsonable_encoder(id)}/cancel",
+            method="GET",
+            request_options=request_options,
+        )
+        try:
+            if 200 <= _response.status_code < 300:
+                return AsyncHttpResponse(response=_response, data=None)
+            if _response.status_code == 501:
+                raise NotImplementedError(
+                    headers=dict(_response.headers),
+                    body=typing.cast(
+                        typing.Optional[typing.Any],
+                        parse_obj_as(
+                            type_=typing.Optional[typing.Any],  # type: ignore
+                            object_=_response.json(),
+                        ),
+                    ),
+                )
+            _response_json = _response.json()
+        except JSONDecodeError:
+            raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
     async def find_optimizations(
         self,
         *,
@@ -371,6 +526,7 @@ class AsyncRawOptimizationsClient:
         dataset_id: typing.Optional[str] = None,
         name: typing.Optional[str] = None,
         dataset_deleted: typing.Optional[bool] = None,
+        filters: typing.Optional[str] = None,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> AsyncHttpResponse[OptimizationPagePublic]:
         """
@@ -388,6 +544,8 @@ class AsyncRawOptimizationsClient:
 
         dataset_deleted : typing.Optional[bool]
 
+        filters : typing.Optional[str]
+
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
@@ -405,6 +563,7 @@ class AsyncRawOptimizationsClient:
                 "dataset_id": dataset_id,
                 "name": name,
                 "dataset_deleted": dataset_deleted,
+                "filters": filters,
             },
             request_options=request_options,
         )
@@ -442,7 +601,8 @@ class AsyncRawOptimizationsClient:
         status: OptimizationWriteStatus,
         id: typing.Optional[str] = OMIT,
         name: typing.Optional[str] = OMIT,
-        metadata: typing.Optional[JsonNodeWrite] = OMIT,
+        metadata: typing.Optional[JsonListStringWrite] = OMIT,
+        studio_config: typing.Optional[OptimizationStudioConfigWrite] = OMIT,
         last_updated_at: typing.Optional[dt.datetime] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> AsyncHttpResponse[None]:
@@ -461,7 +621,9 @@ class AsyncRawOptimizationsClient:
 
         name : typing.Optional[str]
 
-        metadata : typing.Optional[JsonNodeWrite]
+        metadata : typing.Optional[JsonListStringWrite]
+
+        studio_config : typing.Optional[OptimizationStudioConfigWrite]
 
         last_updated_at : typing.Optional[dt.datetime]
 
@@ -481,7 +643,12 @@ class AsyncRawOptimizationsClient:
                 "dataset_name": dataset_name,
                 "objective_name": objective_name,
                 "status": status,
-                "metadata": metadata,
+                "metadata": convert_and_respect_annotation_metadata(
+                    object_=metadata, annotation=JsonListStringWrite, direction="write"
+                ),
+                "studio_config": convert_and_respect_annotation_metadata(
+                    object_=studio_config, annotation=OptimizationStudioConfigWrite, direction="write"
+                ),
                 "last_updated_at": last_updated_at,
             },
             headers={
@@ -506,7 +673,8 @@ class AsyncRawOptimizationsClient:
         status: OptimizationWriteStatus,
         id: typing.Optional[str] = OMIT,
         name: typing.Optional[str] = OMIT,
-        metadata: typing.Optional[JsonNodeWrite] = OMIT,
+        metadata: typing.Optional[JsonListStringWrite] = OMIT,
+        studio_config: typing.Optional[OptimizationStudioConfigWrite] = OMIT,
         last_updated_at: typing.Optional[dt.datetime] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> AsyncHttpResponse[None]:
@@ -525,7 +693,9 @@ class AsyncRawOptimizationsClient:
 
         name : typing.Optional[str]
 
-        metadata : typing.Optional[JsonNodeWrite]
+        metadata : typing.Optional[JsonListStringWrite]
+
+        studio_config : typing.Optional[OptimizationStudioConfigWrite]
 
         last_updated_at : typing.Optional[dt.datetime]
 
@@ -545,7 +715,12 @@ class AsyncRawOptimizationsClient:
                 "dataset_name": dataset_name,
                 "objective_name": objective_name,
                 "status": status,
-                "metadata": metadata,
+                "metadata": convert_and_respect_annotation_metadata(
+                    object_=metadata, annotation=JsonListStringWrite, direction="write"
+                ),
+                "studio_config": convert_and_respect_annotation_metadata(
+                    object_=studio_config, annotation=OptimizationStudioConfigWrite, direction="write"
+                ),
                 "last_updated_at": last_updated_at,
             },
             headers={
@@ -694,3 +869,52 @@ class AsyncRawOptimizationsClient:
         except JSONDecodeError:
             raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
         raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+    async def get_studio_optimization_logs(
+        self, id: str, *, request_options: typing.Optional[RequestOptions] = None
+    ) -> AsyncHttpResponse[OptimizationStudioLog]:
+        """
+        Get presigned S3 URL for downloading optimization logs
+
+        Parameters
+        ----------
+        id : str
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        AsyncHttpResponse[OptimizationStudioLog]
+            Logs response
+        """
+        _response = await self._client_wrapper.httpx_client.request(
+            f"v1/private/optimizations/studio/{jsonable_encoder(id)}/logs",
+            method="GET",
+            request_options=request_options,
+        )
+        try:
+            if 200 <= _response.status_code < 300:
+                _data = typing.cast(
+                    OptimizationStudioLog,
+                    parse_obj_as(
+                        type_=OptimizationStudioLog,  # type: ignore
+                        object_=_response.json(),
+                    ),
+                )
+                return AsyncHttpResponse(response=_response, data=_data)
+            if _response.status_code == 404:
+                raise NotFoundError(
+                    headers=dict(_response.headers),
+                    body=typing.cast(
+                        typing.Optional[typing.Any],
+                        parse_obj_as(
+                            type_=typing.Optional[typing.Any],  # type: ignore
+                            object_=_response.json(),
+                        ),
+                    ),
+                )
+            _response_json = _response.json()
+        except JSONDecodeError:
+            raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)

opik/rest_api/optimizations/types/optimization_update_status.py

@@ -2,4 +2,6 @@
 
 import typing
 
-OptimizationUpdateStatus = typing.Union[typing.Literal["running", "completed", "cancelled"], typing.Any]
+OptimizationUpdateStatus = typing.Union[
+    typing.Literal["running", "completed", "cancelled", "initialized", "error"], typing.Any
+]

opik/rest_api/prompts/__init__.py

@@ -2,6 +2,6 @@
 
 # isort: skip_file
 
-from .types import PromptWriteType
+from .types import CreatePromptVersionDetailTemplateStructure, PromptWriteTemplateStructure, PromptWriteType
 
-__all__ = ["PromptWriteType"]
+__all__ = ["CreatePromptVersionDetailTemplateStructure", "PromptWriteTemplateStructure", "PromptWriteType"]