vibe-bot/vibe_bot/config.py

"""Configuration module for the vibe bot."""

from __future__ import annotations

import logging
import os

from dotenv import load_dotenv

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

load_dotenv()

# Discord
DISCORD_TOKEN: str = os.getenv("DISCORD_TOKEN", "")

# Endpoints
CHAT_ENDPOINT: str = os.getenv("CHAT_ENDPOINT", "")
COMPLETION_ENDPOINT: str = os.getenv("COMPLETION_ENDPOINT", "")
IMAGE_GEN_ENDPOINT: str = os.getenv("IMAGE_GEN_ENDPOINT", "")
IMAGE_EDIT_ENDPOINT: str = os.getenv("IMAGE_EDIT_ENDPOINT", "")
EMBEDDING_ENDPOINT: str = os.getenv("EMBEDDING_ENDPOINT", "")
MAX_COMPLETION_TOKENS: int = int(os.getenv("MAX_COMPLETION_TOKENS", "1000"))

# API Keys
CHAT_ENDPOINT_KEY: str = os.getenv("CHAT_ENDPOINT_KEY", "placeholder")
COMPLETION_ENDPOINT_KEY: str = os.getenv("COMPLETION_ENDPOINT_KEY", "placeholder")
IMAGE_GEN_ENDPOINT_KEY: str = os.getenv("IMAGE_GEN_ENDPOINT_KEY", "placeholder")
IMAGE_EDIT_ENDPOINT_KEY: str = os.getenv("IMAGE_EDIT_ENDPOINT_KEY", "placeholder")
EMBEDDING_ENDPOINT_KEY: str = os.getenv("EMBEDDING_ENDPOINT_KEY", "placeholder")

# Models
CHAT_MODEL: str = os.getenv("CHAT_MODEL", "")
COMPLETION_MODEL: str = os.getenv("COMPLETION_MODEL", "")
IMAGE_GEN_MODEL: str = os.getenv("IMAGE_GEN_MODEL", "")
IMAGE_EDIT_MODEL: str = os.getenv("IMAGE_EDIT_MODEL", "")
EMBEDDING_MODEL: str = os.getenv("EMBEDDING_MODEL", "")

# Database and embeddings
DB_PATH: str = os.getenv("DB_PATH", "chat_history.db")
EMBEDDING_DIMENSION: int = 2048
MAX_HISTORY_MESSAGES: int = int(os.getenv("MAX_HISTORY_MESSAGES", "1000"))
SIMILARITY_THRESHOLD: float = float(os.getenv("SIMILARITY_THRESHOLD", "0.7"))
TOP_K_RESULTS: int = int(os.getenv("TOP_K_RESULTS", "5"))

# Check token
if not DISCORD_TOKEN:
    msg = "DISCORD_TOKEN required."
    raise RuntimeError(msg)

# Check endpoints
if not CHAT_ENDPOINT:
    endpoint_msg = "CHAT_ENDPOINT required."
    raise RuntimeError(endpoint_msg)

if not COMPLETION_ENDPOINT:
    endpoint_msg = "COMPLETION_ENDPOINT required."
    raise RuntimeError(endpoint_msg)

if not IMAGE_GEN_ENDPOINT:
    endpoint_msg = "IMAGE_GEN_ENDPOINT required."
    raise RuntimeError(endpoint_msg)

if not IMAGE_EDIT_ENDPOINT:
    endpoint_msg = "IMAGE_EDIT_ENDPOINT required."
    raise RuntimeError(endpoint_msg)

if not EMBEDDING_ENDPOINT:
    endpoint_msg = "EMBEDDING_ENDPOINT required."
    raise RuntimeError(endpoint_msg)

# Check models
if not CHAT_MODEL:
    model_msg = "CHAT_MODEL required."
    raise RuntimeError(model_msg)

if not COMPLETION_MODEL:
    model_msg = "COMPLETION_MODEL required."
    raise RuntimeError(model_msg)

if not IMAGE_GEN_MODEL:
    model_msg = "IMAGE_GEN_MODEL required."
    raise RuntimeError(model_msg)

if not IMAGE_EDIT_MODEL:
    model_msg = "IMAGE_EDIT_MODEL required."
    raise RuntimeError(model_msg)

if not EMBEDDING_MODEL:
    model_msg = "EMBEDDING_MODEL required."
    raise RuntimeError(model_msg)

# TTS
TTS_MODEL_PATH: str = os.getenv("TTS_MODEL_PATH", "kokoro-v1.0.onnx")
TTS_VOICES_PATH: str = os.getenv("TTS_VOICES_PATH", "voices-v1.0.bin")
TTS_VOICE: str = os.getenv("TTS_VOICE", "af_sarah")
TTS_SPEED: float = float(os.getenv("TTS_SPEED", "1.0"))

# Available voices organized by category
VOICES_LIST: dict[str, dict[str, str | list[str]]] = {
    "🇺🇸 👩": {
        "language": "en-us",
        "voices": [
            "af_alloy",
            "af_aoede",
            "af_bella",
            "af_heart",
            "af_jessica",
            "af_kore",
            "af_nicole",
            "af_nova",
            "af_river",
            "af_sarah",
            "af_sky",
        ],
    },
    "🇺🇸 👨": {
        "language": "en-us",
        "voices": [
            "am_adam",
            "am_echo",
            "am_eric",
            "am_fenrir",
            "am_liam",
            "am_michael",
            "am_onyx",
            "am_puck",
        ],
    },
    "🇬🇧": {
        "language": "en-gb",
        "voices": [
            "bf_alice",
            "bf_emma",
            "bf_isabella",
            "bf_lily",
            "bm_daniel",
            "bm_fable",
            "bm_george",
            "bm_lewis",
        ],
    },
    "🇫🇷": {
        "language": "fr-fr",
        "voices": ["ff_siwis"],
    },
    "🇮🇹": {
        "language": "it",
        "voices": ["if_sara", "im_nicola"],
    },
    "🇯🇵": {
        "language": "ja",
        "voices": [
            "jf_alpha",
            "jf_gongitsune",
            "jf_nezumi",
            "jf_tebukuro",
            "jm_kumo",
        ],
    },
    "🇨🇳": {
        "language": "cmn",
        "voices": [
            "zf_xiaobei",
            "zf_xiaoni",
            "zf_xiaoxiao",
            "zf_xiaoyi",
            "zm_yunjian",
            "zm_yunxi",
            "zm_yunxia",
            "zm_yunyang",
        ],
    },
}

logger.info("CHAT_ENDPOINT set to %s", CHAT_ENDPOINT)
logger.info("COMPLETION_ENDPOINT set to %s", COMPLETION_ENDPOINT)
logger.info("IMAGE_GEN_ENDPOINT set to %s", IMAGE_GEN_ENDPOINT)
logger.info("IMAGE_EDIT_ENDPOINT set to %s", IMAGE_EDIT_ENDPOINT)
logger.info("EMBEDDING_ENDPOINT set to %s", EMBEDDING_ENDPOINT)