# coding: utf-8

"""
    Mixpeek API

    This is the Mixpeek API, providing access to various endpoints for data processing and retrieval.

    The version of the OpenAPI document: 1.3.0
    Contact: info@mixpeek.com
    Generated by OpenAPI Generator (https://openapi-generator.tech)

    Do not edit the class manually.
"""  # noqa: E501


from __future__ import annotations
import pprint
import re  # noqa: F401
import json

from pydantic import BaseModel, ConfigDict, Field, StrictBool, StrictStr, field_validator
from typing import Any, ClassVar, Dict, List, Optional
from typing_extensions import Annotated
from typing import Optional, Set
from typing_extensions import Self

class StageParamsQueryExpand(BaseModel):
    """
    Parameters for query expansion stage.  This stage: 1. Takes your original query 2. Uses an LLM to generate semantically similar query variations 3. Executes feature_search for each variation (original + expansions) 4. Fuses results using Reciprocal Rank Fusion (RRF)  Example - Basic query expansion (copy and run):     ```python     {         \"stages\": [             {                 \"stage\": \"query_expand\",                 \"parameters\": {                     \"num_expansions\": 3,                     \"feature_search_config\": {                         \"query\": {\"text\": \"{{INPUT.query}}\"},                         \"feature_extractors\": [                             {\"field_name\": \"content.text\", \"embedding_model\": \"text\"}                         ],                         \"top_k\": 10                     }                 }             }         ]     }     ```  Example - With custom expansion prompt:     ```python     {         \"stages\": [             {                 \"stage\": \"query_expand\",                 \"parameters\": {                     \"num_expansions\": 5,                     \"expansion_prompt\": \"Generate {{NUM_EXPANSIONS}} alternative search queries for: {{QUERY}}. Focus on synonyms and related concepts. Return one query per line.\",                     \"feature_search_config\": {                         \"query\": {\"text\": \"{{INPUT.query}}\"},                         \"feature_extractors\": [                             {\"field_name\": \"content.text\", \"embedding_model\": \"text\"}                         ],                         \"top_k\": 20                     },                     \"rrf_k\": 60                 }             }         ]     }  Example - Multimodal query expansion:     ```python     {         \"stages\": [             {                 \"stage\": \"query_expand\",                 \"parameters\": {                     \"num_expansions\": 3,                     \"feature_search_config\": {                         \"query\": {\"text\": \"{{INPUT.query}}\", \"image\": \"{{INPUT.image_url}}\"},                         \"feature_extractors\": [                             {\"field_name\": \"content.text\", \"embedding_model\": \"text\"},                             {\"field_name\": \"content.image\", \"embedding_model\": \"multimodal\"}                         ],                         \"top_k\": 15                     },                     \"include_original\": true                 }             }         ]     }  How it works:     1. The original query (from feature_search_config.query.text) is sent to an LLM     2. LLM generates `num_expansions` alternative queries     3. feature_search runs for original query + each expansion     4. Results are fused using RRF: score = sum(1 / (k + rank)) across all queries     5. Documents appearing in multiple result sets get boosted  Why use query expansion:     - Handles vocabulary mismatch (user says \"car\", docs say \"vehicle\")     - Captures related concepts the user might not have thought of     - Improves recall without sacrificing precision (RRF handles fusion)     - Works with any feature_search configuration (text, image, multimodal)
    """ # noqa: E501
    num_expansions: Optional[Annotated[int, Field(le=10, strict=True, ge=1)]] = Field(default=3, description="Number of query variations to generate. More expansions = better recall but slower.")
    expansion_prompt: Optional[StrictStr] = Field(default='null', description="Custom prompt for query expansion. Use {{QUERY}} for the original query and {{NUM_EXPANSIONS}} for the count. If not provided, uses a default prompt.")
    expansion_model: Optional[StrictStr] = Field(default='gpt-4o-mini', description="LLM model to use for generating query expansions.")
    feature_search_config: Dict[str, Any] = Field(description="Full feature_search configuration. This is the same config you would pass to a standalone feature_search stage. The query.text field will be replaced with each expanded query.")
    include_original: Optional[StrictBool] = Field(default=True, description="Whether to include the original query in addition to expansions.")
    rrf_k: Optional[Annotated[int, Field(le=1000, strict=True, ge=1)]] = Field(default=60, description="RRF constant k. Higher values give more weight to lower-ranked results. Default of 60 is standard. Use lower (20-40) for precision, higher (80-100) for recall.")
    fusion_strategy: Optional[StrictStr] = Field(default='rrf', description="How to fuse results from multiple queries. 'rrf' = Reciprocal Rank Fusion (recommended), 'linear' = simple score averaging.")
    deduplicate: Optional[StrictBool] = Field(default=True, description="Whether to deduplicate results by document_id before returning.")
    __properties: ClassVar[List[str]] = ["num_expansions", "expansion_prompt", "expansion_model", "feature_search_config", "include_original", "rrf_k", "fusion_strategy", "deduplicate"]

    @field_validator('fusion_strategy')
    def fusion_strategy_validate_enum(cls, value):
        """Validates the enum"""
        if value is None:
            return value

        if value not in set(['rrf', 'linear']):
            raise ValueError("must be one of enum values ('rrf', 'linear')")
        return value

    model_config = ConfigDict(
        populate_by_name=True,
        validate_assignment=True,
        protected_namespaces=(),
    )


    def to_str(self) -> str:
        """Returns the string representation of the model using alias"""
        return pprint.pformat(self.model_dump(by_alias=True))

    def to_json(self) -> str:
        """Returns the JSON representation of the model using alias"""
        # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
        return json.dumps(self.to_dict())

    @classmethod
    def from_json(cls, json_str: str) -> Optional[Self]:
        """Create an instance of StageParamsQueryExpand from a JSON string"""
        return cls.from_dict(json.loads(json_str))

    def to_dict(self) -> Dict[str, Any]:
        """Return the dictionary representation of the model using alias.

        This has the following differences from calling pydantic's
        `self.model_dump(by_alias=True)`:

        * `None` is only added to the output dict for nullable fields that
          were set at model initialization. Other fields with value `None`
          are ignored.
        """
        excluded_fields: Set[str] = set([
        ])

        _dict = self.model_dump(
            by_alias=True,
            exclude=excluded_fields,
            exclude_none=True,
        )
        return _dict

    @classmethod
    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
        """Create an instance of StageParamsQueryExpand from a dict"""
        if obj is None:
            return None

        if not isinstance(obj, dict):
            return cls.model_validate(obj)

        _obj = cls.model_validate({
            "num_expansions": obj.get("num_expansions") if obj.get("num_expansions") is not None else 3,
            "expansion_prompt": obj.get("expansion_prompt") if obj.get("expansion_prompt") is not None else 'null',
            "expansion_model": obj.get("expansion_model") if obj.get("expansion_model") is not None else 'gpt-4o-mini',
            "feature_search_config": obj.get("feature_search_config"),
            "include_original": obj.get("include_original") if obj.get("include_original") is not None else True,
            "rrf_k": obj.get("rrf_k") if obj.get("rrf_k") is not None else 60,
            "fusion_strategy": obj.get("fusion_strategy") if obj.get("fusion_strategy") is not None else 'rrf',
            "deduplicate": obj.get("deduplicate") if obj.get("deduplicate") is not None else True
        })
        return _obj


