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

import typing
from json.decoder import JSONDecodeError

from ...core.api_error import ApiError
from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ...core.jsonable_encoder import jsonable_encoder
from ...core.pydantic_utilities import pydantic_v1
from ...core.request_options import RequestOptions
from ...types.hosted_by_enum import HostedByEnum
from ...types.ml_model_developer import MlModelDeveloper
from ...types.ml_model_display_config_request import MlModelDisplayConfigRequest
from ...types.ml_model_exec_config_request import MlModelExecConfigRequest
from ...types.ml_model_family import MlModelFamily
from ...types.ml_model_parameter_config_request import MlModelParameterConfigRequest
from ...types.ml_model_read import MlModelRead
from ...types.paginated_ml_model_read_list import PaginatedMlModelReadList
from ...types.visibility_enum import VisibilityEnum

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


class MlModelsClient:
    def __init__(self, *, client_wrapper: SyncClientWrapper):
        self._client_wrapper = client_wrapper

    def list(
        self,
        *,
        limit: typing.Optional[int] = None,
        offset: typing.Optional[int] = None,
        ordering: typing.Optional[str] = None,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> PaginatedMlModelReadList:
        """
        List all ML Models that your Workspace has access to.

        Parameters
        ----------
        limit : typing.Optional[int]
            Number of results to return per page.

        offset : typing.Optional[int]
            The initial index from which to return the results.

        ordering : typing.Optional[str]
            Which field to use when ordering the results.

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

        Returns
        -------
        PaginatedMlModelReadList


        Examples
        --------
        from vellum.client import Vellum

        client = Vellum(
            api_key="YOUR_API_KEY",
        )
        client.ml_models.list()
        """
        _response = self._client_wrapper.httpx_client.request(
            "v1/ml-models",
            base_url=self._client_wrapper.get_environment().default,
            method="GET",
            params={"limit": limit, "offset": offset, "ordering": ordering},
            request_options=request_options,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(PaginatedMlModelReadList, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    def create(
        self,
        *,
        name: str,
        family: MlModelFamily,
        exec_config: MlModelExecConfigRequest,
        hosted_by: typing.Optional[HostedByEnum] = OMIT,
        developed_by: typing.Optional[MlModelDeveloper] = OMIT,
        parameter_config: typing.Optional[MlModelParameterConfigRequest] = OMIT,
        display_config: typing.Optional[MlModelDisplayConfigRequest] = OMIT,
        visibility: typing.Optional[VisibilityEnum] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> MlModelRead:
        """
        Creates a new ML Model.

        Parameters
        ----------
        name : str
            The unique name of the ML Model.

        family : MlModelFamily
            The family of the ML Model.

            * `CAPYBARA` - Capybara
            * `CHAT_GPT` - Chat GPT
            * `CLAUDE` - Claude
            * `COHERE` - Cohere
            * `FALCON` - Falcon
            * `GEMINI` - Gemini
            * `GRANITE` - Granite
            * `GPT3` - GPT-3
            * `FIREWORKS` - Fireworks
            * `LLAMA2` - Llama2
            * `LLAMA3` - Llama3
            * `MISTRAL` - Mistral
            * `MPT` - MPT
            * `OPENCHAT` - OpenChat
            * `PALM` - PaLM
            * `SOLAR` - Solar
            * `TITAN` - Titan
            * `WIZARD` - Wizard
            * `YI` - Yi
            * `ZEPHYR` - Zephyr

        exec_config : MlModelExecConfigRequest
            Configuration for how to execute the ML Model.

        hosted_by : typing.Optional[HostedByEnum]
            The organization hosting the ML Model.

            * `ANTHROPIC` - ANTHROPIC
            * `AWS_BEDROCK` - AWS_BEDROCK
            * `AZURE_OPENAI` - AZURE_OPENAI
            * `COHERE` - COHERE
            * `CUSTOM` - CUSTOM
            * `FIREWORKS_AI` - FIREWORKS_AI
            * `GOOGLE` - GOOGLE
            * `GOOGLE_VERTEX_AI` - GOOGLE_VERTEX_AI
            * `GROQ` - GROQ
            * `HUGGINGFACE` - HUGGINGFACE
            * `IBM_WATSONX` - IBM_WATSONX
            * `MOSAICML` - MOSAICML
            * `MYSTIC` - MYSTIC
            * `OPENAI` - OPENAI
            * `OPENPIPE` - OPENPIPE
            * `PYQ` - PYQ
            * `REPLICATE` - REPLICATE

        developed_by : typing.Optional[MlModelDeveloper]
            The organization that developed the ML Model.

            * `01_AI` - 01_AI
            * `AMAZON` - AMAZON
            * `ANTHROPIC` - ANTHROPIC
            * `COHERE` - COHERE
            * `ELUTHERAI` - ELUTHERAI
            * `FIREWORKS_AI` - FIREWORKS_AI
            * `GOOGLE` - GOOGLE
            * `HUGGINGFACE` - HUGGINGFACE
            * `IBM` - IBM
            * `META` - META
            * `MISTRAL_AI` - MISTRAL_AI
            * `MOSAICML` - MOSAICML
            * `NOUS_RESEARCH` - NOUS_RESEARCH
            * `OPENAI` - OPENAI
            * `OPENCHAT` - OPENCHAT
            * `OPENPIPE` - OPENPIPE
            * `TII` - TII
            * `WIZARDLM` - WIZARDLM

        parameter_config : typing.Optional[MlModelParameterConfigRequest]
            Configuration for the ML Model's parameters.

        display_config : typing.Optional[MlModelDisplayConfigRequest]
            Configuration for how to display the ML Model.

        visibility : typing.Optional[VisibilityEnum]
            The visibility of the ML Model.

            * `DEFAULT` - DEFAULT
            * `PUBLIC` - PUBLIC
            * `PRIVATE` - PRIVATE
            * `DISABLED` - DISABLED

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

        Returns
        -------
        MlModelRead


        Examples
        --------
        from vellum import MlModelExecConfigRequest
        from vellum.client import Vellum

        client = Vellum(
            api_key="YOUR_API_KEY",
        )
        client.ml_models.create(
            name="name",
            family="CAPYBARA",
            exec_config=MlModelExecConfigRequest(
                model_identifier="model_identifier",
                base_url="base_url",
                metadata={"key": "value"},
                features=["TEXT"],
            ),
        )
        """
        _response = self._client_wrapper.httpx_client.request(
            "v1/ml-models",
            base_url=self._client_wrapper.get_environment().default,
            method="POST",
            json={
                "name": name,
                "family": family,
                "hosted_by": hosted_by,
                "developed_by": developed_by,
                "exec_config": exec_config,
                "parameter_config": parameter_config,
                "display_config": display_config,
                "visibility": visibility,
            },
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(MlModelRead, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    def retrieve(self, id: str, *, request_options: typing.Optional[RequestOptions] = None) -> MlModelRead:
        """
        Retrieve an ML Model by its UUID.

        Parameters
        ----------
        id : str
            Either the ML Model's ID or its unique name

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

        Returns
        -------
        MlModelRead


        Examples
        --------
        from vellum.client import Vellum

        client = Vellum(
            api_key="YOUR_API_KEY",
        )
        client.ml_models.retrieve(
            id="id",
        )
        """
        _response = self._client_wrapper.httpx_client.request(
            f"v1/ml-models/{jsonable_encoder(id)}",
            base_url=self._client_wrapper.get_environment().default,
            method="GET",
            request_options=request_options,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(MlModelRead, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    def update(
        self,
        id: str,
        *,
        display_config: typing.Optional[MlModelDisplayConfigRequest] = OMIT,
        visibility: typing.Optional[VisibilityEnum] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> MlModelRead:
        """
        Replace an ML Model with a new representation, keying off of its UUID.

        Parameters
        ----------
        id : str
            Either the ML Model's ID or its unique name

        display_config : typing.Optional[MlModelDisplayConfigRequest]
            Configuration for how to display the ML Model.

        visibility : typing.Optional[VisibilityEnum]
            The visibility of the ML Model.

            * `DEFAULT` - DEFAULT
            * `PUBLIC` - PUBLIC
            * `PRIVATE` - PRIVATE
            * `DISABLED` - DISABLED

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

        Returns
        -------
        MlModelRead


        Examples
        --------
        from vellum.client import Vellum

        client = Vellum(
            api_key="YOUR_API_KEY",
        )
        client.ml_models.update(
            id="id",
        )
        """
        _response = self._client_wrapper.httpx_client.request(
            f"v1/ml-models/{jsonable_encoder(id)}",
            base_url=self._client_wrapper.get_environment().default,
            method="PUT",
            json={"display_config": display_config, "visibility": visibility},
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(MlModelRead, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    def partial_update(
        self,
        id: str,
        *,
        display_config: typing.Optional[MlModelDisplayConfigRequest] = OMIT,
        visibility: typing.Optional[VisibilityEnum] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> MlModelRead:
        """
        Partially update an ML Model, keying off of its UUID.

        Parameters
        ----------
        id : str
            Either the ML Model's ID or its unique name

        display_config : typing.Optional[MlModelDisplayConfigRequest]
            Configuration for how to display the ML Model.

        visibility : typing.Optional[VisibilityEnum]
            The visibility of the ML Model.

            * `DEFAULT` - DEFAULT
            * `PUBLIC` - PUBLIC
            * `PRIVATE` - PRIVATE
            * `DISABLED` - DISABLED

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

        Returns
        -------
        MlModelRead


        Examples
        --------
        from vellum.client import Vellum

        client = Vellum(
            api_key="YOUR_API_KEY",
        )
        client.ml_models.partial_update(
            id="id",
        )
        """
        _response = self._client_wrapper.httpx_client.request(
            f"v1/ml-models/{jsonable_encoder(id)}",
            base_url=self._client_wrapper.get_environment().default,
            method="PATCH",
            json={"display_config": display_config, "visibility": visibility},
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(MlModelRead, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)


class AsyncMlModelsClient:
    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        self._client_wrapper = client_wrapper

    async def list(
        self,
        *,
        limit: typing.Optional[int] = None,
        offset: typing.Optional[int] = None,
        ordering: typing.Optional[str] = None,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> PaginatedMlModelReadList:
        """
        List all ML Models that your Workspace has access to.

        Parameters
        ----------
        limit : typing.Optional[int]
            Number of results to return per page.

        offset : typing.Optional[int]
            The initial index from which to return the results.

        ordering : typing.Optional[str]
            Which field to use when ordering the results.

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

        Returns
        -------
        PaginatedMlModelReadList


        Examples
        --------
        import asyncio

        from vellum.client import AsyncVellum

        client = AsyncVellum(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.ml_models.list()


        asyncio.run(main())
        """
        _response = await self._client_wrapper.httpx_client.request(
            "v1/ml-models",
            base_url=self._client_wrapper.get_environment().default,
            method="GET",
            params={"limit": limit, "offset": offset, "ordering": ordering},
            request_options=request_options,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(PaginatedMlModelReadList, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    async def create(
        self,
        *,
        name: str,
        family: MlModelFamily,
        exec_config: MlModelExecConfigRequest,
        hosted_by: typing.Optional[HostedByEnum] = OMIT,
        developed_by: typing.Optional[MlModelDeveloper] = OMIT,
        parameter_config: typing.Optional[MlModelParameterConfigRequest] = OMIT,
        display_config: typing.Optional[MlModelDisplayConfigRequest] = OMIT,
        visibility: typing.Optional[VisibilityEnum] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> MlModelRead:
        """
        Creates a new ML Model.

        Parameters
        ----------
        name : str
            The unique name of the ML Model.

        family : MlModelFamily
            The family of the ML Model.

            * `CAPYBARA` - Capybara
            * `CHAT_GPT` - Chat GPT
            * `CLAUDE` - Claude
            * `COHERE` - Cohere
            * `FALCON` - Falcon
            * `GEMINI` - Gemini
            * `GRANITE` - Granite
            * `GPT3` - GPT-3
            * `FIREWORKS` - Fireworks
            * `LLAMA2` - Llama2
            * `LLAMA3` - Llama3
            * `MISTRAL` - Mistral
            * `MPT` - MPT
            * `OPENCHAT` - OpenChat
            * `PALM` - PaLM
            * `SOLAR` - Solar
            * `TITAN` - Titan
            * `WIZARD` - Wizard
            * `YI` - Yi
            * `ZEPHYR` - Zephyr

        exec_config : MlModelExecConfigRequest
            Configuration for how to execute the ML Model.

        hosted_by : typing.Optional[HostedByEnum]
            The organization hosting the ML Model.

            * `ANTHROPIC` - ANTHROPIC
            * `AWS_BEDROCK` - AWS_BEDROCK
            * `AZURE_OPENAI` - AZURE_OPENAI
            * `COHERE` - COHERE
            * `CUSTOM` - CUSTOM
            * `FIREWORKS_AI` - FIREWORKS_AI
            * `GOOGLE` - GOOGLE
            * `GOOGLE_VERTEX_AI` - GOOGLE_VERTEX_AI
            * `GROQ` - GROQ
            * `HUGGINGFACE` - HUGGINGFACE
            * `IBM_WATSONX` - IBM_WATSONX
            * `MOSAICML` - MOSAICML
            * `MYSTIC` - MYSTIC
            * `OPENAI` - OPENAI
            * `OPENPIPE` - OPENPIPE
            * `PYQ` - PYQ
            * `REPLICATE` - REPLICATE

        developed_by : typing.Optional[MlModelDeveloper]
            The organization that developed the ML Model.

            * `01_AI` - 01_AI
            * `AMAZON` - AMAZON
            * `ANTHROPIC` - ANTHROPIC
            * `COHERE` - COHERE
            * `ELUTHERAI` - ELUTHERAI
            * `FIREWORKS_AI` - FIREWORKS_AI
            * `GOOGLE` - GOOGLE
            * `HUGGINGFACE` - HUGGINGFACE
            * `IBM` - IBM
            * `META` - META
            * `MISTRAL_AI` - MISTRAL_AI
            * `MOSAICML` - MOSAICML
            * `NOUS_RESEARCH` - NOUS_RESEARCH
            * `OPENAI` - OPENAI
            * `OPENCHAT` - OPENCHAT
            * `OPENPIPE` - OPENPIPE
            * `TII` - TII
            * `WIZARDLM` - WIZARDLM

        parameter_config : typing.Optional[MlModelParameterConfigRequest]
            Configuration for the ML Model's parameters.

        display_config : typing.Optional[MlModelDisplayConfigRequest]
            Configuration for how to display the ML Model.

        visibility : typing.Optional[VisibilityEnum]
            The visibility of the ML Model.

            * `DEFAULT` - DEFAULT
            * `PUBLIC` - PUBLIC
            * `PRIVATE` - PRIVATE
            * `DISABLED` - DISABLED

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

        Returns
        -------
        MlModelRead


        Examples
        --------
        import asyncio

        from vellum import MlModelExecConfigRequest
        from vellum.client import AsyncVellum

        client = AsyncVellum(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.ml_models.create(
                name="name",
                family="CAPYBARA",
                exec_config=MlModelExecConfigRequest(
                    model_identifier="model_identifier",
                    base_url="base_url",
                    metadata={"key": "value"},
                    features=["TEXT"],
                ),
            )


        asyncio.run(main())
        """
        _response = await self._client_wrapper.httpx_client.request(
            "v1/ml-models",
            base_url=self._client_wrapper.get_environment().default,
            method="POST",
            json={
                "name": name,
                "family": family,
                "hosted_by": hosted_by,
                "developed_by": developed_by,
                "exec_config": exec_config,
                "parameter_config": parameter_config,
                "display_config": display_config,
                "visibility": visibility,
            },
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(MlModelRead, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    async def retrieve(self, id: str, *, request_options: typing.Optional[RequestOptions] = None) -> MlModelRead:
        """
        Retrieve an ML Model by its UUID.

        Parameters
        ----------
        id : str
            Either the ML Model's ID or its unique name

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

        Returns
        -------
        MlModelRead


        Examples
        --------
        import asyncio

        from vellum.client import AsyncVellum

        client = AsyncVellum(
            api_key="YOUR_API_KEY",
        )


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


        asyncio.run(main())
        """
        _response = await self._client_wrapper.httpx_client.request(
            f"v1/ml-models/{jsonable_encoder(id)}",
            base_url=self._client_wrapper.get_environment().default,
            method="GET",
            request_options=request_options,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(MlModelRead, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    async def update(
        self,
        id: str,
        *,
        display_config: typing.Optional[MlModelDisplayConfigRequest] = OMIT,
        visibility: typing.Optional[VisibilityEnum] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> MlModelRead:
        """
        Replace an ML Model with a new representation, keying off of its UUID.

        Parameters
        ----------
        id : str
            Either the ML Model's ID or its unique name

        display_config : typing.Optional[MlModelDisplayConfigRequest]
            Configuration for how to display the ML Model.

        visibility : typing.Optional[VisibilityEnum]
            The visibility of the ML Model.

            * `DEFAULT` - DEFAULT
            * `PUBLIC` - PUBLIC
            * `PRIVATE` - PRIVATE
            * `DISABLED` - DISABLED

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

        Returns
        -------
        MlModelRead


        Examples
        --------
        import asyncio

        from vellum.client import AsyncVellum

        client = AsyncVellum(
            api_key="YOUR_API_KEY",
        )


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


        asyncio.run(main())
        """
        _response = await self._client_wrapper.httpx_client.request(
            f"v1/ml-models/{jsonable_encoder(id)}",
            base_url=self._client_wrapper.get_environment().default,
            method="PUT",
            json={"display_config": display_config, "visibility": visibility},
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(MlModelRead, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    async def partial_update(
        self,
        id: str,
        *,
        display_config: typing.Optional[MlModelDisplayConfigRequest] = OMIT,
        visibility: typing.Optional[VisibilityEnum] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> MlModelRead:
        """
        Partially update an ML Model, keying off of its UUID.

        Parameters
        ----------
        id : str
            Either the ML Model's ID or its unique name

        display_config : typing.Optional[MlModelDisplayConfigRequest]
            Configuration for how to display the ML Model.

        visibility : typing.Optional[VisibilityEnum]
            The visibility of the ML Model.

            * `DEFAULT` - DEFAULT
            * `PUBLIC` - PUBLIC
            * `PRIVATE` - PRIVATE
            * `DISABLED` - DISABLED

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

        Returns
        -------
        MlModelRead


        Examples
        --------
        import asyncio

        from vellum.client import AsyncVellum

        client = AsyncVellum(
            api_key="YOUR_API_KEY",
        )


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


        asyncio.run(main())
        """
        _response = await self._client_wrapper.httpx_client.request(
            f"v1/ml-models/{jsonable_encoder(id)}",
            base_url=self._client_wrapper.get_environment().default,
            method="PATCH",
            json={"display_config": display_config, "visibility": visibility},
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                return pydantic_v1.parse_obj_as(MlModelRead, _response.json())  # type: ignore
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)
