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

from __future__ import annotations
from ..core.pydantic_utilities import UniversalBaseModel
from .metadata_filters import MetadataFilters
import typing
import pydantic
from .retrieval_mode import RetrievalMode
from ..core.pydantic_utilities import IS_PYDANTIC_V2
from ..core.pydantic_utilities import update_forward_refs


class PresetRetrievalParams(UniversalBaseModel):
    """
    Schema for the search params for an retrieval execution that can be preset for a pipeline.
    """

    dense_similarity_top_k: typing.Optional[int] = pydantic.Field(default=None)
    """
    Number of nodes for dense retrieval.
    """

    dense_similarity_cutoff: typing.Optional[float] = pydantic.Field(default=None)
    """
    Minimum similarity score wrt query for retrieval
    """

    sparse_similarity_top_k: typing.Optional[int] = pydantic.Field(default=None)
    """
    Number of nodes for sparse retrieval.
    """

    enable_reranking: typing.Optional[bool] = pydantic.Field(default=None)
    """
    Enable reranking for retrieval
    """

    rerank_top_n: typing.Optional[int] = pydantic.Field(default=None)
    """
    Number of reranked nodes for returning.
    """

    alpha: typing.Optional[float] = pydantic.Field(default=None)
    """
    Alpha value for hybrid retrieval to determine the weights between dense and sparse retrieval. 0 is sparse retrieval and 1 is dense retrieval.
    """

    search_filters: typing.Optional[MetadataFilters] = pydantic.Field(default=None)
    """
    Search filters for retrieval.
    """

    files_top_k: typing.Optional[int] = pydantic.Field(default=None)
    """
    Number of files to retrieve (only for retrieval mode files_via_metadata and files_via_content).
    """

    retrieval_mode: typing.Optional[RetrievalMode] = pydantic.Field(default=None)
    """
    The retrieval mode for the query.
    """

    retrieve_image_nodes: typing.Optional[bool] = pydantic.Field(default=None)
    """
    Whether to retrieve image nodes.
    """

    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


update_forward_refs(MetadataFilters, PresetRetrievalParams=PresetRetrievalParams)
