from pydantic.dataclasses import dataclass
import os

@dataclass
class LLM:
    name: str  # On Azure OpenAI Service, this is the deployment name
    max_tokens: int


# name of the model deployment on Azure OpenAI Service
@dataclass
class GPT_4_OMNI(LLM):
    name: str = "gpt-4o-2" if os.environ.get("OPENAI_API_TYPE") == "azure" else "gpt-4o"
    max_tokens: int = 128000


@dataclass
class GPT_4_128k(LLM):
    name: str = "gpt-4-1106-preview"
    max_tokens: int = 128000


@dataclass
class GPT_4_32k(LLM):
    name: str = "gpt4-32k"
    max_tokens: int = 32768


@dataclass
class GPT_4(LLM):
    name: str = "gpt-4-0613"
    max_tokens: int = 8192


@dataclass
class GPT_3_5_Turbo_16k(LLM):
    name: str = "gpt-3.5-turbo-16k"
    max_tokens: int = 16384


@dataclass
class GPT_3_5_Turbo(LLM):
    name: str = "gpt-3.5-turbo"
    max_tokens: int = 4096


EXPLAIN_MODEL: LLM = GPT_4_OMNI()
PLAN_MODEL: LLM = GPT_4_OMNI()
EXECUTE_MODEL: LLM = GPT_4_OMNI()
OPTIMIZE_MODEL: LLM = GPT_4_OMNI()
