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

from ..core.pydantic_utilities import UniversalBaseModel
import typing
import pydantic
from ..core.pydantic_utilities import IS_PYDANTIC_V2


class CohereEmbedding(UniversalBaseModel):
    model_name: typing.Optional[str] = pydantic.Field(default=None)
    """
    The modelId of the Cohere model to use.
    """

    embed_batch_size: typing.Optional[int] = pydantic.Field(default=None)
    """
    The batch size for embedding calls.
    """

    num_workers: typing.Optional[int] = pydantic.Field(default=None)
    """
    The number of workers to use for async embedding calls.
    """

    api_key: typing.Optional[str] = pydantic.Field(default=None)
    """
    The Cohere API key.
    """

    truncate: typing.Optional[str] = pydantic.Field(default=None)
    """
    Truncation type - START/ END/ NONE
    """

    input_type: typing.Optional[str] = pydantic.Field(default=None)
    """
    Model Input type. If not provided, search_document and search_query are used when needed.
    """

    embedding_type: typing.Optional[str] = pydantic.Field(default=None)
    """
    Embedding type. If not provided float embedding_type is used when needed.
    """

    class_name: typing.Optional[str] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow
