# 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, StrictStr, field_validator
from typing import Any, ClassVar, Dict, List, Optional
from typing_extensions import Annotated
from mixpeek.models.citation_config import CitationConfig
from typing import Optional, Set
from typing_extensions import Self

class StageParamsRagPrepare(BaseModel):
    """
    Configuration for RAG context preparation.  **Stage Category**: APPLY  **Transformation**: N documents → 1 context document (single_context mode)                    OR N documents → N formatted documents (formatted_list mode)  **Purpose**: Prepare search results for LLM consumption by formatting documents, managing token budgets, and adding citations. This is a preparation stage that does NOT call an LLM - it prepares content for downstream LLM stages.  **When to Use**:     - Before passing search results to an LLM for RAG     - When you need to fit multiple documents into a token budget     - When you need citation tracking for source attribution     - When you need consistent document formatting  **When NOT to Use**:     - When you want the LLM to generate a summary (use summarize stage)     - When you don't need token management     - For simple pass-through of documents  **Output Modes**:     - `single_context`: Combines all documents into one context string     - `formatted_list`: Returns individually formatted documents  **Common Pipeline Position**: feature_search → rerank → rag_prepare → (external LLM call)  Examples:     Basic context preparation:         ```json         {             \"max_tokens\": 8000,             \"output_mode\": \"single_context\"         }         ```      Custom document formatting with citations:         ```json         {             \"max_tokens\": 4000,             \"document_template\": \"[{{CONTEXT.INDEX}}] {{DOC.metadata.title}}\\n{{DOC.content}}\\n\",             \"citation\": {\"style\": \"numbered\", \"include_title\": true}         }         ```      Formatted list for custom processing:         ```json         {             \"output_mode\": \"formatted_list\",             \"document_template\": \"Source: {{DOC.metadata.source}}\\n{{DOC.content}}\"         }         ```
    """ # noqa: E501
    max_tokens: Optional[Annotated[int, Field(le=128000, strict=True, ge=100)]] = Field(default=8000, description="OPTIONAL. Maximum tokens for the combined context output. Documents exceeding this limit are handled by truncation_strategy. Default: 8000 (safe for most models).")
    tokenizer: Optional[StrictStr] = Field(default='cl100k_base', description="OPTIONAL. Tokenizer to use for token counting. Default: 'cl100k_base' (GPT-4/GPT-3.5 tokenizer). Options: 'cl100k_base', 'p50k_base', 'r50k_base', 'gpt2'")
    truncation_strategy: Optional[StrictStr] = Field(default='priority_truncate', description="OPTIONAL. How to handle documents exceeding max_tokens: - 'priority_truncate': Include docs in score order, truncate last to fit - 'proportional': Give each doc proportional token budget based on count - 'drop_last': Include complete docs until limit, drop remaining")
    output_mode: Optional[StrictStr] = Field(default='single_context', description="OPTIONAL. Output format: - 'single_context': One document with combined 'context' string + 'citations' - 'formatted_list': N documents with 'formatted_content' field each")
    document_template: Optional[StrictStr] = Field(default='''[{{CONTEXT.INDEX}}] {{DOC.content}}

''', description="OPTIONAL. Template for formatting each document. Available placeholders: - {{CONTEXT.INDEX}}: 1-based position in result set (1, 2, 3...) - {{CONTEXT.CITATION}}: Citation marker based on citation.style - {{DOC.*}}: Any document field (e.g., {{DOC.content}}, {{DOC.metadata.title}})")
    content_field: Optional[StrictStr] = Field(default='content', description="Primary field to extract content from each document.")
    separator: Optional[StrictStr] = Field(default='''
''', description="Separator between documents in single_context mode.")
    citation: Optional[CitationConfig] = Field(default=None, description="Citation configuration for source tracking.")
    context_field: Optional[StrictStr] = Field(default='rag_context', description="Field name for combined context (single_context mode).")
    citations_field: Optional[StrictStr] = Field(default='citations', description="Field name for citation metadata.")
    formatted_content_field: Optional[StrictStr] = Field(default='formatted_content', description="Field name for formatted content (formatted_list mode).")
    __properties: ClassVar[List[str]] = ["max_tokens", "tokenizer", "truncation_strategy", "output_mode", "document_template", "content_field", "separator", "citation", "context_field", "citations_field", "formatted_content_field"]

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

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

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

        if value not in set(['single_context', 'formatted_list']):
            raise ValueError("must be one of enum values ('single_context', 'formatted_list')")
        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 StageParamsRagPrepare 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,
        )
        # override the default output from pydantic by calling `to_dict()` of citation
        if self.citation:
            _dict['citation'] = self.citation.to_dict()
        return _dict

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

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

        _obj = cls.model_validate({
            "max_tokens": obj.get("max_tokens") if obj.get("max_tokens") is not None else 8000,
            "tokenizer": obj.get("tokenizer") if obj.get("tokenizer") is not None else 'cl100k_base',
            "truncation_strategy": obj.get("truncation_strategy") if obj.get("truncation_strategy") is not None else 'priority_truncate',
            "output_mode": obj.get("output_mode") if obj.get("output_mode") is not None else 'single_context',
            "document_template": obj.get("document_template") if obj.get("document_template") is not None else '''[{{CONTEXT.INDEX}}] {{DOC.content}}

''',
            "content_field": obj.get("content_field") if obj.get("content_field") is not None else 'content',
            "separator": obj.get("separator") if obj.get("separator") is not None else '''
''',
            "citation": CitationConfig.from_dict(obj["citation"]) if obj.get("citation") is not None else None,
            "context_field": obj.get("context_field") if obj.get("context_field") is not None else 'rag_context',
            "citations_field": obj.get("citations_field") if obj.get("citations_field") is not None else 'citations',
            "formatted_content_field": obj.get("formatted_content_field") if obj.get("formatted_content_field") is not None else 'formatted_content'
        })
        return _obj


