client.py

# This file was auto-generated by Fern from our API Definition.

import typing

from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ..core.request_options import RequestOptions
from .raw_client import AsyncRawBatchesClient, RawBatchesClient
from .types.batch import Batch
from .types.cancel_batch_response import CancelBatchResponse
from .types.create_batch_response import CreateBatchResponse
from .types.get_batch_response import GetBatchResponse
from .types.list_batches_response import ListBatchesResponse

# this is used as the default value for optional parameters
OMIT = typing.cast(typing.Any, ...)


class BatchesClient:
    def __init__(self, *, client_wrapper: SyncClientWrapper):
        self._raw_client = RawBatchesClient(client_wrapper=client_wrapper)

    @property
    def with_raw_response(self) -> RawBatchesClient:
        """
        Retrieves a raw implementation of this client that returns raw responses.

        Returns
        -------
        RawBatchesClient
        """
        return self._raw_client

    def list(
        self,
        *,
        page_size: typing.Optional[int] = None,
        page_token: typing.Optional[str] = None,
        order_by: typing.Optional[str] = None,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> ListBatchesResponse:
        """
        List the batches for the current user

        Parameters
        ----------
        page_size : typing.Optional[int]
            The maximum number of batches to return. The service may return fewer than
            this value.
            If unspecified, at most 50 batches will be returned.
            The maximum value is 1000; values above 1000 will be coerced to 1000.

        page_token : typing.Optional[str]
            A page token, received from a previous `ListBatches` call.
            Provide this to retrieve the subsequent page.

        order_by : typing.Optional[str]
            Batches can be ordered by creation time or last updated time.
            Use `created_at` for creation time or `updated_at` for last updated time.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        ListBatchesResponse
            A successful response.

        Examples
        --------
        from cohere import Client

        client = Client(
            client_name="YOUR_CLIENT_NAME",
            token="YOUR_TOKEN",
        )
        client.batches.list(
            page_size=1,
            page_token="page_token",
            order_by="order_by",
        )
        """
        _response = self._raw_client.list(
            page_size=page_size, page_token=page_token, order_by=order_by, request_options=request_options
        )
        return _response.data

    def create(self, *, request: Batch, request_options: typing.Optional[RequestOptions] = None) -> CreateBatchResponse:
        """
        Creates and executes a batch from an uploaded dataset of requests

        Parameters
        ----------
        request : Batch

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        CreateBatchResponse
            A successful response.

        Examples
        --------
        from cohere import Client
        from cohere.batches import Batch

        client = Client(
            client_name="YOUR_CLIENT_NAME",
            token="YOUR_TOKEN",
        )
        client.batches.create(
            request=Batch(
                name="name",
                input_dataset_id="input_dataset_id",
                model="model",
            ),
        )
        """
        _response = self._raw_client.create(request=request, request_options=request_options)
        return _response.data

    def retrieve(self, id: str, *, request_options: typing.Optional[RequestOptions] = None) -> GetBatchResponse:
        """
        Retrieves a batch

        Parameters
        ----------
        id : str
            The batch ID.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        GetBatchResponse
            A successful response.

        Examples
        --------
        from cohere import Client

        client = Client(
            client_name="YOUR_CLIENT_NAME",
            token="YOUR_TOKEN",
        )
        client.batches.retrieve(
            id="id",
        )
        """
        _response = self._raw_client.retrieve(id, request_options=request_options)
        return _response.data

    def cancel(self, id: str, *, request_options: typing.Optional[RequestOptions] = None) -> CancelBatchResponse:
        """
        Cancels an in-progress batch

        Parameters
        ----------
        id : str
            The batch ID.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        CancelBatchResponse
            A successful response.

        Examples
        --------
        from cohere import Client

        client = Client(
            client_name="YOUR_CLIENT_NAME",
            token="YOUR_TOKEN",
        )
        client.batches.cancel(
            id="id",
        )
        """
        _response = self._raw_client.cancel(id, request_options=request_options)
        return _response.data


class AsyncBatchesClient:
    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        self._raw_client = AsyncRawBatchesClient(client_wrapper=client_wrapper)

    @property
    def with_raw_response(self) -> AsyncRawBatchesClient:
        """
        Retrieves a raw implementation of this client that returns raw responses.

        Returns
        -------
        AsyncRawBatchesClient
        """
        return self._raw_client

    async def list(
        self,
        *,
        page_size: typing.Optional[int] = None,
        page_token: typing.Optional[str] = None,
        order_by: typing.Optional[str] = None,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> ListBatchesResponse:
        """
        List the batches for the current user

        Parameters
        ----------
        page_size : typing.Optional[int]
            The maximum number of batches to return. The service may return fewer than
            this value.
            If unspecified, at most 50 batches will be returned.
            The maximum value is 1000; values above 1000 will be coerced to 1000.

        page_token : typing.Optional[str]
            A page token, received from a previous `ListBatches` call.
            Provide this to retrieve the subsequent page.

        order_by : typing.Optional[str]
            Batches can be ordered by creation time or last updated time.
            Use `created_at` for creation time or `updated_at` for last updated time.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        ListBatchesResponse
            A successful response.

        Examples
        --------
        import asyncio

        from cohere import AsyncClient

        client = AsyncClient(
            client_name="YOUR_CLIENT_NAME",
            token="YOUR_TOKEN",
        )


        async def main() -> None:
            await client.batches.list(
                page_size=1,
                page_token="page_token",
                order_by="order_by",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.list(
            page_size=page_size, page_token=page_token, order_by=order_by, request_options=request_options
        )
        return _response.data

    async def create(
        self, *, request: Batch, request_options: typing.Optional[RequestOptions] = None
    ) -> CreateBatchResponse:
        """
        Creates and executes a batch from an uploaded dataset of requests

        Parameters
        ----------
        request : Batch

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        CreateBatchResponse
            A successful response.

        Examples
        --------
        import asyncio

        from cohere import AsyncClient
        from cohere.batches import Batch

        client = AsyncClient(
            client_name="YOUR_CLIENT_NAME",
            token="YOUR_TOKEN",
        )


        async def main() -> None:
            await client.batches.create(
                request=Batch(
                    name="name",
                    input_dataset_id="input_dataset_id",
                    model="model",
                ),
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.create(request=request, request_options=request_options)
        return _response.data

    async def retrieve(self, id: str, *, request_options: typing.Optional[RequestOptions] = None) -> GetBatchResponse:
        """
        Retrieves a batch

        Parameters
        ----------
        id : str
            The batch ID.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        GetBatchResponse
            A successful response.

        Examples
        --------
        import asyncio

        from cohere import AsyncClient

        client = AsyncClient(
            client_name="YOUR_CLIENT_NAME",
            token="YOUR_TOKEN",
        )


        async def main() -> None:
            await client.batches.retrieve(
                id="id",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.retrieve(id, request_options=request_options)
        return _response.data

    async def cancel(self, id: str, *, request_options: typing.Optional[RequestOptions] = None) -> CancelBatchResponse:
        """
        Cancels an in-progress batch

        Parameters
        ----------
        id : str
            The batch ID.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        CancelBatchResponse
            A successful response.

        Examples
        --------
        import asyncio

        from cohere import AsyncClient

        client = AsyncClient(
            client_name="YOUR_CLIENT_NAME",
            token="YOUR_TOKEN",
        )


        async def main() -> None:
            await client.batches.cancel(
                id="id",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.cancel(id, request_options=request_options)
        return _response.data