"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

from __future__ import annotations
from .logicaloperator_input import LogicalOperatorInput, LogicalOperatorInputTypedDict
from mixpeek.types import BaseModel, Nullable, OptionalNullable, UNSET, UNSET_SENTINEL
from pydantic import model_serializer
from typing import Any, Dict, Optional
from typing_extensions import NotRequired, TypedDict


class StageConfigInputTypedDict(TypedDict):
    r"""Configuration for how a retriever stage processes queries and handles results.

    This model consolidates stage-specific settings that control behavior
    during retrieval operations, making it easier to configure stages
    with consistent options.
    """

    stage_id: str
    r"""Unique identifier for the stage"""
    stage_name: str
    r"""Name of the stage"""
    version: str
    r"""Version of the stage"""
    parameters: NotRequired[Dict[str, Any]]
    r"""Parameters for the stage"""
    pre_filters: NotRequired[Nullable[LogicalOperatorInputTypedDict]]
    r"""Filters to apply before the main search is executed"""
    post_filters: NotRequired[Nullable[LogicalOperatorInputTypedDict]]
    r"""Filters to apply to search results after retrieval"""
    timeout_seconds: NotRequired[Nullable[float]]
    r"""Maximum execution time for this stage in seconds"""
    exclude_from_final: NotRequired[bool]
    r"""If True, results from this stage are used for subsequent stages but excluded from final output"""


class StageConfigInput(BaseModel):
    r"""Configuration for how a retriever stage processes queries and handles results.

    This model consolidates stage-specific settings that control behavior
    during retrieval operations, making it easier to configure stages
    with consistent options.
    """

    stage_id: str
    r"""Unique identifier for the stage"""

    stage_name: str
    r"""Name of the stage"""

    version: str
    r"""Version of the stage"""

    parameters: Optional[Dict[str, Any]] = None
    r"""Parameters for the stage"""

    pre_filters: OptionalNullable[LogicalOperatorInput] = UNSET
    r"""Filters to apply before the main search is executed"""

    post_filters: OptionalNullable[LogicalOperatorInput] = UNSET
    r"""Filters to apply to search results after retrieval"""

    timeout_seconds: OptionalNullable[float] = UNSET
    r"""Maximum execution time for this stage in seconds"""

    exclude_from_final: Optional[bool] = False
    r"""If True, results from this stage are used for subsequent stages but excluded from final output"""

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = [
            "parameters",
            "pre_filters",
            "post_filters",
            "timeout_seconds",
            "exclude_from_final",
        ]
        nullable_fields = ["pre_filters", "post_filters", "timeout_seconds"]
        null_default_fields = []

        serialized = handler(self)

        m = {}

        for n, f in type(self).model_fields.items():
            k = f.alias or n
            val = serialized.get(k)
            serialized.pop(k, None)

            optional_nullable = k in optional_fields and k in nullable_fields
            is_set = (
                self.__pydantic_fields_set__.intersection({n})
                or k in null_default_fields
            )  # pylint: disable=no-member

            if val is not None and val != UNSET_SENTINEL:
                m[k] = val
            elif val != UNSET_SENTINEL and (
                not k in optional_fields or (optional_nullable and is_set)
            ):
                m[k] = val

        return m
