# Code generated by Lark OpenAPI.

from typing import Any, Optional, Union, Dict, List, Set, IO, Callable, Type
from lark_oapi.core.construct import init
from .llm_message import LlmMessage


class LlmConfig(object):
    _types = {
        "model": str,
        "messages": List[LlmMessage],
        "max_tokens": int,
        "message_type": str,
        "n": int,
        "temperature": float,
        "presence_penalty": float,
        "frequency_penalty": float,
    }

    def __init__(self, d=None):
        self.model: Optional[str] = None
        self.messages: Optional[List[LlmMessage]] = None
        self.max_tokens: Optional[int] = None
        self.message_type: Optional[str] = None
        self.n: Optional[int] = None
        self.temperature: Optional[float] = None
        self.presence_penalty: Optional[float] = None
        self.frequency_penalty: Optional[float] = None
        init(self, d, self._types)

    @staticmethod
    def builder() -> "LlmConfigBuilder":
        return LlmConfigBuilder()


class LlmConfigBuilder(object):
    def __init__(self) -> None:
        self._llm_config = LlmConfig()

    def model(self, model: str) -> "LlmConfigBuilder":
        self._llm_config.model = model
        return self

    def messages(self, messages: List[LlmMessage]) -> "LlmConfigBuilder":
        self._llm_config.messages = messages
        return self

    def max_tokens(self, max_tokens: int) -> "LlmConfigBuilder":
        self._llm_config.max_tokens = max_tokens
        return self

    def message_type(self, message_type: str) -> "LlmConfigBuilder":
        self._llm_config.message_type = message_type
        return self

    def n(self, n: int) -> "LlmConfigBuilder":
        self._llm_config.n = n
        return self

    def temperature(self, temperature: float) -> "LlmConfigBuilder":
        self._llm_config.temperature = temperature
        return self

    def presence_penalty(self, presence_penalty: float) -> "LlmConfigBuilder":
        self._llm_config.presence_penalty = presence_penalty
        return self

    def frequency_penalty(self, frequency_penalty: float) -> "LlmConfigBuilder":
        self._llm_config.frequency_penalty = frequency_penalty
        return self

    def build(self) -> "LlmConfig":
        return self._llm_config