fix linting, formatting, and add tests

This commit is contained in:
2026-05-23 19:06:53 -04:00
parent 5e708b009c
commit 6ec9fbe85f
15 changed files with 2911 additions and 491 deletions
+62 -45
View File
@@ -1,91 +1,108 @@
from dotenv import load_dotenv
import os
"""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"
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
load_dotenv()
# Discord
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN", "")
DISCORD_TOKEN: str = os.getenv("DISCORD_TOKEN", "")
# Endpoints
CHAT_ENDPOINT = os.getenv("CHAT_ENDPOINT", "")
COMPLETION_ENDPOINT = os.getenv("COMPLETION_ENDPOINT", "")
IMAGE_GEN_ENDPOINT = os.getenv("IMAGE_GEN_ENDPOINT", "")
IMAGE_EDIT_ENDPOINT = os.getenv("IMAGE_EDIT_ENDPOINT", "")
EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "")
MAX_COMPLETION_TOKENS = int(os.getenv("MAX_COMPLETION_TOKENS", "1000"))
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 = os.getenv("CHAT_ENDPOINT_KEY", "placeholder")
COMPLETION_ENDPOINT_KEY = os.getenv("COMPLETION_ENDPOINT_KEY", "placeholder")
IMAGE_GEN_ENDPOINT_KEY = os.getenv("IMAGE_GEN_ENDPOINT_KEY", "placeholder")
IMAGE_EDIT_ENDPOINT_KEY = os.getenv("IMAGE_EDIT_ENDPOINT_KEY", "placeholder")
EMBEDDING_ENDPOINT_KEY = os.getenv("EMBEDDING_ENDPOINT_KEY", "placeholder")
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 = os.getenv("CHAT_MODEL", "")
COMPLETION_MODEL = os.getenv("COMPLETION_MODEL", "")
IMAGE_GEN_MODEL = os.getenv("IMAGE_GEN_MODEL", "")
IMAGE_EDIT_MODEL = os.getenv("IMAGE_EDIT_MODEL", "")
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "")
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 = os.getenv("DB_PATH", "chat_history.db")
EMBEDDING_DIMENSION = 2048
MAX_HISTORY_MESSAGES = int(os.getenv("MAX_HISTORY_MESSAGES", "1000"))
SIMILARITY_THRESHOLD = float(os.getenv("SIMILARITY_THRESHOLD", "0.7"))
TOP_K_RESULTS = int(os.getenv("TOP_K_RESULTS", "5"))
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:
raise Exception("DISCORD_TOKEN required.")
msg = "DISCORD_TOKEN required."
raise RuntimeError(msg)
# Check endpoints
if not CHAT_ENDPOINT:
raise Exception("CHAT_ENDPOINT required.")
endpoint_msg = "CHAT_ENDPOINT required."
raise RuntimeError(endpoint_msg)
if not COMPLETION_ENDPOINT:
raise Exception("COMPLETION_ENDPOINT required.")
endpoint_msg = "COMPLETION_ENDPOINT required."
raise RuntimeError(endpoint_msg)
if not IMAGE_GEN_ENDPOINT:
raise Exception("IMAGE_GEN_ENDPOINT required.")
endpoint_msg = "IMAGE_GEN_ENDPOINT required."
raise RuntimeError(endpoint_msg)
if not IMAGE_EDIT_ENDPOINT:
raise Exception("IMAGE_EDIT_ENDPOINT required.")
endpoint_msg = "IMAGE_EDIT_ENDPOINT required."
raise RuntimeError(endpoint_msg)
if not EMBEDDING_ENDPOINT:
raise Exception("EMBEDDING_ENDPOINT required.")
endpoint_msg = "EMBEDDING_ENDPOINT required."
raise RuntimeError(endpoint_msg)
# Check models
if not CHAT_MODEL:
raise Exception("CHAT_MODEL required.")
model_msg = "CHAT_MODEL required."
raise RuntimeError(model_msg)
if not COMPLETION_MODEL:
raise Exception("COMPLETION_MODEL required.")
model_msg = "COMPLETION_MODEL required."
raise RuntimeError(model_msg)
if not IMAGE_GEN_MODEL:
raise Exception("IMAGE_GEN_MODEL required.")
model_msg = "IMAGE_GEN_MODEL required."
raise RuntimeError(model_msg)
if not IMAGE_EDIT_MODEL:
raise Exception("IMAGE_EDIT_MODEL required.")
model_msg = "IMAGE_EDIT_MODEL required."
raise RuntimeError(model_msg)
if not EMBEDDING_MODEL:
raise Exception("EMBEDDING_MODEL required.")
model_msg = "EMBEDDING_MODEL required."
raise RuntimeError(model_msg)
# TTS
TTS_MODEL_PATH = os.getenv("TTS_MODEL_PATH", "kokoro-v1.0.onnx")
TTS_VOICES_PATH = os.getenv("TTS_VOICES_PATH", "voices-v1.0.bin")
TTS_VOICE = os.getenv("TTS_VOICE", "af_sarah")
TTS_SPEED = float(os.getenv("TTS_SPEED", "1.0"))
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"))
logger.info(f"CHAT_ENDPOINT set to {CHAT_ENDPOINT}")
logger.info(f"COMPLETION_ENDPOINT set to {COMPLETION_ENDPOINT}")
logger.info(f"IMAGE_GEN_ENDPOINT set to {IMAGE_GEN_ENDPOINT}")
logger.info(f"IMAGE_EDIT_ENDPOINT set to {IMAGE_EDIT_ENDPOINT}")
logger.info(f"EMBEDDING_ENDPOINT set to {EMBEDDING_ENDPOINT}")
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)