# Standard Import
from enum import Enum


class LLM(str, Enum):
    GPT4 = "gpt4"
    Claude = "claude"
    Gemini = "gemini"


class ModelType(str, Enum):
    LogReg = "log_reg"
    DT = "dt"
    KNN = "knn"


class Task(str, Enum):
    Verbalization = "verbalization"
    Generation = "generation"


class Dataset(str, Enum):
    Blood = "blood"
    Car = "car"
    Diabetes = "diabetes"


# Open AI Stuff
OPENAI_MODEL_NAMES = {
    LLM.GPT4: "gpt-4o",
    # LLM.GPT4: "gpt-4-turbo",
}

# Claude Config
CLAUDE_LOCATION = "us-east5"
CLAUDE_MODEL_NAMES = {
    LLM.Claude: "claude-3-5-sonnet@20240620"
}

# Gemini Config
GOOGLE_LOCATION = "us-east4"
GOOGLE_MODEL_NAMES = {
    LLM.Gemini: "gemini-1.5-pro-001",
}

GOOGLE_VERBALIZATION_CONFIG = {
    "temperature": 0.1,
    "top_p": 0.95,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

GOOGLE_GENERATION_CONFIG = {
    "temperature": 0.1,
    "top_p": 0.95,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}