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

import datetime as dt
import typing

from ..core.datetime_utils import serialize_datetime
from .base_prompt_template import BasePromptTemplate
from .pooling import Pooling
from .pydantic_program_mode import PydanticProgramMode

try:
    import pydantic.v1 as pydantic  # type: ignore
except ImportError:
    import pydantic  # type: ignore


class HuggingFaceInferenceApiEmbedding(pydantic.BaseModel):
    """
    Wrapper on the Hugging Face's Inference API for embeddings.

    Overview of the design:

    - Uses the feature extraction task: https://huggingface.co/tasks/feature-extraction
    """

    model_name: typing.Optional[str] = pydantic.Field(
        description="The model to run inference with. Can be a model id or URL."
    )
    embed_batch_size: typing.Optional[int] = pydantic.Field(description="The batch size for embedding calls.")
    callback_manager: typing.Optional[typing.Dict[str, typing.Any]]
    system_prompt: typing.Optional[str] = pydantic.Field(description="System prompt for LLM calls.")
    output_parser: typing.Optional[typing.Dict[str, typing.Any]] = pydantic.Field(
        description="Output parser to parse, validate, and correct errors programmatically."
    )
    pydantic_program_mode: typing.Optional[PydanticProgramMode]
    query_wrapper_prompt: typing.Optional[BasePromptTemplate] = pydantic.Field(
        description="Query wrapper prompt for LLM calls."
    )
    token: typing.Optional[str] = pydantic.Field(description="Hugging Face API token.")
    timeout: typing.Optional[float] = pydantic.Field(
        description="The maximum number of seconds to wait for a response from the server. Loading a new model in Inference API can take up to several minutes. Defaults to None, meaning it will loop until the server is available."
    )
    headers: typing.Optional[typing.Dict[str, str]] = pydantic.Field(
        description="Additional headers to send to the server. By default only the authorization and user-agent headers are sent. Values in this dictionary will override the default values."
    )
    cookies: typing.Optional[typing.Dict[str, str]] = pydantic.Field(
        description="Additional cookies to send to the server."
    )
    task: typing.Optional[str] = pydantic.Field(
        description="Optional task to pick Hugging Face's recommended model, used when model_name is left as default of None."
    )
    context_window: typing.Optional[int] = pydantic.Field(
        description="Total number of tokens the model can be input and output for one response. This may be looked up in a model's `config.json`."
    )
    num_output: typing.Optional[int] = pydantic.Field(
        description="Number of tokens the model can output when generating a response."
    )
    is_chat_model: typing.Optional[bool] = pydantic.Field(
        description="Set True if the model exposes a chat interface (i.e. can be passed a sequence of messages, rather than text), like OpenAI's /v1/chat/completions endpoint. Unless chat templating is intentionally applied, Hugging Face models are not chat models."
    )
    is_function_calling_model: typing.Optional[bool] = pydantic.Field(
        description="Set True if the model supports function calling messages, similar to OpenAI's function calling API. For example, converting 'Email Anya to see if she wants to get coffee next Friday' to a function call like `send_email(to: string, body: string)`. As of 10/17/2023, Hugging Face doesn't support function calling messages."
    )
    pooling: typing.Optional[Pooling] = pydantic.Field(
        description="Pooling technique to use with embeddings capability."
    )
    query_instruction: typing.Optional[str] = pydantic.Field(
        description="Instruction to prepend during query embedding."
    )
    text_instruction: typing.Optional[str] = pydantic.Field(description="Instruction to prepend during text embedding.")
    class_name: typing.Optional[str]

    def json(self, **kwargs: typing.Any) -> str:
        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
        return super().json(**kwargs_with_defaults)

    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
        return super().dict(**kwargs_with_defaults)

    class Config:
        frozen = True
        smart_union = True
        json_encoders = {dt.datetime: serialize_datetime}
