# 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
from typing import Any, ClassVar, Dict, List, Optional, Union
from typing_extensions import Annotated
from mixpeek.models.llm_provider import LLMProvider
from mixpeek.models.logical_operator import LogicalOperator
from typing import Optional, Set
from typing_extensions import Self

class StageParamsLlmEnrich(BaseModel):
    """
    Configuration for augmenting documents with LLM generated fields.  **Stage Category**: APPLY (1-1 Enrichment/Generation)  **Transformation**: N documents → N documents (same count, expanded schema)  **Purpose**: Applies LLM generation to each input document, creating new fields with generated content. Use this for summaries, insights, descriptions, or transformations. Each input document produces exactly one output document with added generated fields.  **When to Use**:     - After FILTER/SORT to enhance final results with generated content     - For summarization of long content     - To extract structured data (entities, insights, key points)     - For content transformation (translation, rephrasing, formatting)     - To generate descriptions, titles, or metadata     - For creative augmentation (suggestions, recommendations)  **When NOT to Use**:     - For removing documents (use FILTER: llm_filter instead)     - For simple field transformations (use direct field mapping)     - For initial document retrieval (use FILTER: hybrid_search)     - For reordering (use SORT stages)     - When fast response time is critical (LLM generation is slow, 200ms-5s)     - When cost is a major concern (LLM generation is very expensive)     - For large batch processing (consider async batch jobs instead)  **Operational Behavior**:     - Applies LLM generation to each input document (1-1 operation)     - Maintains document count: N in → N out     - Expands schema: adds new generated fields to each document     - Makes HTTP requests to Engine service for LLM inference     - Very slow operation (LLM generation, 200ms-5s per document batch)     - Processes documents in batches to optimize throughput     - Supports concurrent batching for parallel LLM calls  **Common Pipeline Position**: FILTER → SORT → APPLY (this stage)  **Cost & Performance**:     - Very Expensive: LLM generation costs per document (10-100x vs embeddings)     - Very Slow: 200ms-5s per batch depending on LLM and generation length     - CRITICAL: Use `when` parameter for selective enrichment (massive cost savings)     - Consider enriching only top-ranked results after RANK stage     - Smaller batch sizes often better for latency  **Conditional Enrichment**: Supports `when` parameter to only enrich specific documents. CRITICAL FOR COST SAVINGS - LLM generation is expensive!     - Only summarize long documents (word_count > 500)     - Only process high-priority items     - Only enrich specific content types (articles, not images)  Requirements:     - provider: OPTIONAL, LLM provider (openai, google, anthropic). Auto-inferred if not specified.     - model_name: OPTIONAL, specific model name. Uses provider default if not specified.     - prompt: REQUIRED, LLM prompt template (supports {DOC.field}, {INPUT.field})     - output_field: REQUIRED, where to store generated content     - batch_size: OPTIONAL, documents per batch (default 5)     - schema: OPTIONAL, JSON schema for structured output     - when: OPTIONAL but RECOMMENDED for cost control  Use Cases:     - Summarization: Generate 3-sentence summaries of articles     - Insight extraction: Extract key takeaways and insights     - Description generation: Create product descriptions from specs     - Translation: Translate content to other languages     - Entity extraction: Extract people, places, organizations     - Recommendation generation: Create personalized suggestions  Examples:     Unconditional enrichment:         ```json         {             \"provider\": \"openai\",             \"model_name\": \"gpt-4o-mini\",             \"prompt\": \"Summarize the document\",             \"output_field\": \"metadata.summary\"         }         ```      Conditional enrichment (only summarize long documents):         ```json         {             \"provider\": \"google\",             \"model_name\": \"gemini-2.0-flash\",             \"prompt\": \"Summarize the document\",             \"output_field\": \"metadata.summary\",             \"when\": {                 \"field\": \"metadata.word_count\",                 \"operator\": \"gt\",                 \"value\": 500             }         }         ```
    """ # noqa: E501
    provider: Optional[LLMProvider] = Field(default=None, description="LLM provider to use. Supported providers: - openai: GPT models (GPT-4o, GPT-4o-mini) - google: Gemini models (Gemini 2.0 Flash) - anthropic: Claude models (Claude 3.5 Sonnet/Haiku)  If not specified, defaults to 'google'. Can be auto-inferred from model_name.")
    model_name: Optional[StrictStr] = Field(default='null', description="Specific LLM model to use. If not specified, uses provider default. Examples: gemini-2.0-flash, gpt-4o-mini, claude-3-5-haiku-20241022")
    inference_name: Optional[StrictStr] = Field(default=None, description="DEPRECATED: Use 'provider' and 'model_name' instead. Legacy format: 'provider:model' (e.g., 'gemini:gemini-2.0-flash'). Kept for backward compatibility only.")
    prompt: Optional[StrictStr] = Field(default='Summarize the following content in 2-3 sentences: {{DOC.content}}', description="Prompt template for the LLM (supports doc/input templates).")
    output_field: Optional[StrictStr] = Field(default='metadata.summary', description="Dot-path where the enrichment result should be stored.")
    batch_size: Optional[Annotated[int, Field(le=25, strict=True, ge=1)]] = Field(default=5, description="Number of documents to enrich per LLM request batch.")
    var_schema: Optional[Dict[str, Any]] = Field(default=None, description="Optional JSON schema instructions for the LLM output.", alias="schema")
    temperature: Optional[Union[Annotated[float, Field(le=1.0, strict=True, ge=0.0)], Annotated[int, Field(le=1, strict=True, ge=0)]]] = Field(default=0.2, description="Sampling temperature passed to the LLM.")
    api_key: Optional[StrictStr] = Field(default=None, description="OPTIONAL. Bring Your Own Key (BYOK) - use your own LLM API key instead of Mixpeek's.  **How to use:** 1. Store your API key as an organization secret via POST /v1/organizations/secrets    Example: {\"secret_name\": \"openai_api_key\", \"secret_value\": \"sk-proj-...\"}  2. Reference it here using template syntax: {{secrets.openai_api_key}}  **Benefits:** - Use your own API credits and rate limits - Keep your API keys secure in Mixpeek's encrypted vault - No changes needed to your retriever when rotating keys  If not provided, uses Mixpeek's default API keys (usage charged to your account).")
    when: Optional[LogicalOperator] = Field(default=None, description="OPTIONAL. Conditional filter that documents must satisfy to be enriched with LLM. Uses LogicalOperator (AND/OR/NOT) for complex boolean logic, or simple field/operator/value for single conditions. Documents NOT matching this condition will SKIP enrichment (pass-through unchanged). CRITICAL FOR COST SAVINGS - LLM calls are expensive! Only enrich documents that need it. When NOT specified, ALL documents are enriched unconditionally (may incur high costs).   Use cases: - Only summarize documents with word_count > 500 - Only enrich English articles/blogs - Only process high-priority items   Simple condition example: {\"field\": \"metadata.word_count\", \"operator\": \"gt\", \"value\": 500} Boolean AND example: {\"AND\": [{\"field\": \"category\", \"operator\": \"in\", \"value\": [\"article\"]}, ...]} ")
    __properties: ClassVar[List[str]] = ["provider", "model_name", "inference_name", "prompt", "output_field", "batch_size", "schema", "temperature", "api_key", "when"]

    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 StageParamsLlmEnrich 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 when
        if self.when:
            _dict['when'] = self.when.to_dict()
        return _dict

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

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

        _obj = cls.model_validate({
            "provider": obj.get("provider"),
            "model_name": obj.get("model_name") if obj.get("model_name") is not None else 'null',
            "inference_name": obj.get("inference_name"),
            "prompt": obj.get("prompt") if obj.get("prompt") is not None else 'Summarize the following content in 2-3 sentences: {{DOC.content}}',
            "output_field": obj.get("output_field") if obj.get("output_field") is not None else 'metadata.summary',
            "batch_size": obj.get("batch_size") if obj.get("batch_size") is not None else 5,
            "schema": obj.get("schema"),
            "temperature": obj.get("temperature") if obj.get("temperature") is not None else 0.2,
            "api_key": obj.get("api_key"),
            "when": LogicalOperator.from_dict(obj["when"]) if obj.get("when") is not None else None
        })
        return _obj


