settings.py

import torch
from loguru import logger
from pydantic_settings import BaseSettings, SettingsConfigDict
from zenml.client import Client
from zenml.exceptions import EntityExistsError


class Settings(BaseSettings):
    model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
    # OpenAI
    OPENAI_API_KEY: str | None = None
    OPENAI_MODEL_ID: str | None = None

    # ollama model
    OLLAMA_MODEL_ID: str | None = None
    # ollama inference params
    # OLLAMA_BASE_URL: str = "http://localhost:11434"

    # huggingface model
    HUGGINGFACE_INFERENCE_MODEL_ID: str | None = None
    TEMPERATURE_INFERENCE: float = 0.15
    MAX_NEW_TOKENS_INFERENCE: int = 1024
    TOP_P_INFERENCE: float = 0.90
    REPETITION_PENALTY: float = 1.05
    HF_MODEL_DEVICE: str = (
        "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    )

    # huggingface token
    HUGGINGFACE_ACCESS_TOKEN: str | None = None

    # Comet ML (during training)
    COMET_API_KEY: str | None = None
    COMET_PROJECT: str | None = None

    # MongoDB database
    DATABASE_HOST: str | None = None  # mongodb://localhost:27017/
    DATABASE_NAME: str | None = None

    # Qdrant vector database
    USE_QDRANT_CLOUD: bool = False
    QDRANT_DATABASE_HOST: str | None = None  # localhost
    QDRANT_DATABASE_PORT: int | None = None  # 6333
    QDRANT_CLOUD_URL: str = "str"
    QDRANT_APIKEY: str | None = None

    # RAG
    TEXT_EMBEDDING_MODEL_ID: str = "sentence-transformers/all-MiniLM-L6-v2"
    RERANKING_CROSS_ENCODER_MODEL_ID: str = "cross-encoder/ms-marco-MiniLM-L4-v2"
    RAG_MODEL_DEVICE: str = (
        "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    )

    # uvicorn
    UVICORN_HOST: str | None = None
    UVICORN_PORT: int = 8000
    UVICORN_RELOAD: bool = True

    @property
    def OPENAI_MAX_TOKEN_WINDOW(self) -> int:
        official_max_token_window = {
            "gpt-4o-mini": 128_000,
            "gpt-4.1-nano": 1_047_576,
        }.get(self.OPENAI_MODEL_ID, 128_000)

        max_token_window = int(official_max_token_window * 0.90)  # Use 90% of the official max token window to be safe

        return max_token_window

    @classmethod
    def load_settings(cls) -> "Settings":
        """
        Load settings from the ZenML secret store or fallback to .env file.
        :return:Settings: The initialized settings object.
        """
        try:
            logger.info("Loading settings from the ZenML secret store.")

            settings_secrets = Client().get_secret("settings")
            loaded_settings = Settings(**settings_secrets.secret_values)
        except (RuntimeError, KeyError):
            logger.warning(
                "Failed to load settings from the ZenML secret store. Defaulting to loading the settings from the '.env' file."
            )
            loaded_settings = Settings()

        return loaded_settings

    def export(self) -> None:
        """
        Exports the settings to the ZenML secret store.
        """

        env_vars = self.model_dump()
        for key, value in env_vars.items():
            env_vars[key] = str(value)

        client = Client()

        try:
            client.create_secret(name="settings", values=env_vars)
        except EntityExistsError:
            logger.warning(
                "Secret 'scope' already exists. Delete it manually by running 'zenml secret delete settings', before trying to recreate it."
            )


settings = Settings.load_settings()