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

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


class VertexTextEmbedding(UniversalBaseModel):
    model_name: typing.Optional[str] = pydantic.Field(default=None)
    """
    The modelId of the VertexAI 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.
    """

    location: str = pydantic.Field()
    """
    The default location to use when making API calls.
    """

    project: str = pydantic.Field()
    """
    The default GCP project to use when making Vertex API calls.
    """

    embed_mode: typing.Optional[VertexEmbeddingMode] = pydantic.Field(default=None)
    """
    The embedding mode to use.
    """

    additional_kwargs: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
    """
    Additional kwargs for the Vertex.
    """

    client_email: typing.Optional[str] = pydantic.Field(default=None)
    """
    The client email for the VertexAI credentials.
    """

    token_uri: typing.Optional[str] = pydantic.Field(default=None)
    """
    The token URI for the VertexAI credentials.
    """

    private_key_id: typing.Optional[str] = pydantic.Field(default=None)
    """
    The private key ID for the VertexAI credentials.
    """

    private_key: typing.Optional[str] = pydantic.Field(default=None)
    """
    The private key for the VertexAI credentials.
    """

    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
