from pydantic import BaseModel, Field


class AiSettingBase(BaseModel):
    provider: str = Field(default="nvidia")
    api_base: str = Field(default="https://integrate.api.nvidia.com/v1")
    api_key: str = Field(default="")
    model_name: str = Field(default="nvidia/llama-3.1-nemotron-70b-instruct")
    temperature: float = Field(default=0.7, ge=0, le=2)
    max_tokens: int = Field(default=2048, ge=1)

    # 提示词
    generate_prompt: str = Field(default="")
    optimizer_quality_prompt: str = Field(default="")
    optimizer_polish_prompt: str = Field(default="")
    optimizer_evaluate_prompt: str = Field(default="")
    crawler_extract_prompt: str = Field(default="")
    crawler_rewrite_prompt: str = Field(default="")

    # 优化器各阶段温度
    optimizer_quality_temperature: float = Field(default=0.3, ge=0, le=2)
    optimizer_polish_temperature: float = Field(default=0.8, ge=0, le=2)
    optimizer_evaluate_temperature: float = Field(default=0.3, ge=0, le=2)

    crawl_enabled: bool = Field(default=False)
    crawl_keywords: str = Field(default="")
    max_pages_per_run: int = Field(default=3, ge=1)


class AiSettingCreate(AiSettingBase):
    pass