refactor: isolate database migrations and model defaults
This commit is contained in:
@@ -42,7 +42,10 @@ def configure_logging() -> None:
|
||||
|
||||
|
||||
def create_services(settings: AppSettings) -> ApplicationServices:
|
||||
storage = AssistantStorage(settings.db_path)
|
||||
storage = AssistantStorage(
|
||||
settings.db_path,
|
||||
default_models=settings.default_models(),
|
||||
)
|
||||
ai_client: AIClient
|
||||
speech_recognizer: SpeechTranscriber
|
||||
if settings.mode == "local":
|
||||
|
||||
@@ -66,6 +66,15 @@ class AppSettings:
|
||||
whisper_compute_type: str
|
||||
whisper_language: str | None
|
||||
|
||||
def default_models(self) -> dict[str, str]:
|
||||
models = {"local": self.ollama_model}
|
||||
if self.yandex_cloud_folder:
|
||||
models["yandex"] = (
|
||||
f"gpt://{self.yandex_cloud_folder}/"
|
||||
f"{self.yandex_cloud_model}"
|
||||
)
|
||||
return models
|
||||
|
||||
|
||||
def load_env_file(path: Path = ENV_FILE) -> None:
|
||||
if not path.exists():
|
||||
@@ -209,154 +218,3 @@ def load_app_settings() -> AppSettings:
|
||||
else whisper_language or None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def get_bot_token() -> str:
|
||||
load_env_file()
|
||||
token = os.getenv(TOKEN_ENV_NAME)
|
||||
if not token:
|
||||
raise RuntimeError(f"Set {TOKEN_ENV_NAME} in environment or .env file.")
|
||||
return token
|
||||
|
||||
|
||||
def get_assistant_password() -> str:
|
||||
load_env_file()
|
||||
password = os.getenv(PASSWORD_ENV_NAME, "").strip()
|
||||
if not password:
|
||||
raise RuntimeError(
|
||||
f"Set {PASSWORD_ENV_NAME} in environment or .env file."
|
||||
)
|
||||
return password
|
||||
|
||||
|
||||
def get_db_path() -> Path:
|
||||
load_env_file()
|
||||
raw_path = os.getenv(DB_ENV_NAME)
|
||||
if not raw_path:
|
||||
return DEFAULT_DB_FILE
|
||||
|
||||
path = Path(raw_path).expanduser()
|
||||
if path.is_absolute():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
def get_assistant_mode() -> str:
|
||||
load_env_file()
|
||||
mode = os.getenv(MODE_ENV_NAME, DEFAULT_MODE).strip().lower()
|
||||
if mode not in {"local", "yandex"}:
|
||||
raise RuntimeError(
|
||||
f"{MODE_ENV_NAME} must be either 'local' or 'yandex', got {mode!r}."
|
||||
)
|
||||
return mode
|
||||
|
||||
|
||||
def get_ollama_base_url() -> str:
|
||||
load_env_file()
|
||||
return os.getenv(OLLAMA_BASE_URL_ENV_NAME, DEFAULT_OLLAMA_BASE_URL).rstrip("/")
|
||||
|
||||
|
||||
def get_default_ollama_model() -> str:
|
||||
load_env_file()
|
||||
return os.getenv(OLLAMA_MODEL_ENV_NAME, DEFAULT_OLLAMA_MODEL)
|
||||
|
||||
|
||||
def get_yandex_cloud_folder() -> str:
|
||||
load_env_file()
|
||||
folder = os.getenv(YANDEX_CLOUD_FOLDER_ENV_NAME, "").strip()
|
||||
if not folder or folder == "":
|
||||
raise RuntimeError(
|
||||
f"Set {YANDEX_CLOUD_FOLDER_ENV_NAME} in environment or .env file."
|
||||
)
|
||||
return folder.strip("/")
|
||||
|
||||
|
||||
def get_yandex_cloud_model() -> str:
|
||||
load_env_file()
|
||||
model = os.getenv(
|
||||
YANDEX_CLOUD_MODEL_ENV_NAME,
|
||||
DEFAULT_YANDEX_CLOUD_MODEL,
|
||||
).strip()
|
||||
return model.strip("/") or DEFAULT_YANDEX_CLOUD_MODEL
|
||||
|
||||
|
||||
def get_default_yandex_model() -> str:
|
||||
return f"gpt://{get_yandex_cloud_folder()}/{get_yandex_cloud_model()}"
|
||||
|
||||
|
||||
def get_default_model(provider: str) -> str:
|
||||
if provider == "local":
|
||||
return get_default_ollama_model()
|
||||
if provider == "yandex":
|
||||
return get_default_yandex_model()
|
||||
raise ValueError(f"Unknown AI provider: {provider}")
|
||||
|
||||
|
||||
def get_yandex_stt_model() -> str:
|
||||
load_env_file()
|
||||
return os.getenv(YANDEX_STT_MODEL_ENV_NAME, DEFAULT_YANDEX_STT_MODEL).strip()
|
||||
|
||||
|
||||
def get_yandex_stt_language() -> str:
|
||||
load_env_file()
|
||||
return os.getenv(
|
||||
YANDEX_STT_LANGUAGE_ENV_NAME,
|
||||
DEFAULT_YANDEX_STT_LANGUAGE,
|
||||
).strip()
|
||||
|
||||
|
||||
def get_local_timezone() -> ZoneInfo:
|
||||
load_env_file()
|
||||
timezone_name = os.getenv(TIMEZONE_ENV_NAME, DEFAULT_TIMEZONE)
|
||||
try:
|
||||
return ZoneInfo(timezone_name)
|
||||
except ZoneInfoNotFoundError:
|
||||
logger.warning("Unknown timezone %s, falling back to UTC", timezone_name)
|
||||
return ZoneInfo("UTC")
|
||||
|
||||
|
||||
def get_whisper_model() -> str:
|
||||
load_env_file()
|
||||
return os.getenv(WHISPER_MODEL_ENV_NAME, DEFAULT_WHISPER_MODEL).strip()
|
||||
|
||||
|
||||
def get_whisper_device() -> str:
|
||||
load_env_file()
|
||||
return os.getenv(WHISPER_DEVICE_ENV_NAME, DEFAULT_WHISPER_DEVICE).strip()
|
||||
|
||||
|
||||
def get_whisper_compute_type() -> str:
|
||||
load_env_file()
|
||||
return os.getenv(
|
||||
WHISPER_COMPUTE_TYPE_ENV_NAME,
|
||||
DEFAULT_WHISPER_COMPUTE_TYPE,
|
||||
).strip()
|
||||
|
||||
|
||||
def get_whisper_language() -> str | None:
|
||||
load_env_file()
|
||||
language = os.getenv(
|
||||
WHISPER_LANGUAGE_ENV_NAME,
|
||||
DEFAULT_WHISPER_LANGUAGE,
|
||||
).strip()
|
||||
return None if language.lower() == "auto" else language or None
|
||||
|
||||
|
||||
def get_voice_max_duration_seconds() -> int:
|
||||
load_env_file()
|
||||
raw_value = os.getenv(
|
||||
VOICE_MAX_DURATION_ENV_NAME,
|
||||
str(DEFAULT_VOICE_MAX_DURATION_SECONDS),
|
||||
)
|
||||
try:
|
||||
value = int(raw_value)
|
||||
except ValueError:
|
||||
value = 0
|
||||
if value <= 0:
|
||||
logger.warning(
|
||||
"%s must be a positive integer, using %s",
|
||||
VOICE_MAX_DURATION_ENV_NAME,
|
||||
DEFAULT_VOICE_MAX_DURATION_SECONDS,
|
||||
)
|
||||
return DEFAULT_VOICE_MAX_DURATION_SECONDS
|
||||
return value
|
||||
|
||||
113
assistant_bot/migrations.py
Normal file
113
assistant_bot/migrations.py
Normal file
@@ -0,0 +1,113 @@
|
||||
import sqlite3
|
||||
|
||||
|
||||
LATEST_SCHEMA_VERSION = 2
|
||||
|
||||
INITIAL_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS user_settings (
|
||||
user_id INTEGER PRIMARY KEY,
|
||||
ollama_model TEXT,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS authorized_users (
|
||||
user_id INTEGER PRIMARY KEY,
|
||||
authorized_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS memories (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS notes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS reminders (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
chat_id INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
remind_at TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
created_at TEXT NOT NULL,
|
||||
sent_at TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tracked_items (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversation_messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
chat_id INTEGER NOT NULL,
|
||||
role TEXT NOT NULL CHECK(role IN ('user', 'assistant')),
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversation_contexts (
|
||||
user_id INTEGER NOT NULL,
|
||||
chat_id INTEGER NOT NULL,
|
||||
started_after_id INTEGER NOT NULL DEFAULT 0,
|
||||
updated_at TEXT NOT NULL,
|
||||
PRIMARY KEY (user_id, chat_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_memories_user_created
|
||||
ON memories(user_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_notes_user_created
|
||||
ON notes(user_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_reminders_due
|
||||
ON reminders(status, remind_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_tracked_user_updated
|
||||
ON tracked_items(user_id, updated_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_chat_id
|
||||
ON conversation_messages(user_id, chat_id, id DESC);
|
||||
"""
|
||||
|
||||
|
||||
def _schema_version(connection: sqlite3.Connection) -> int:
|
||||
row = connection.execute("PRAGMA user_version").fetchone()
|
||||
return int(row[0])
|
||||
|
||||
|
||||
def _has_column(
|
||||
connection: sqlite3.Connection,
|
||||
table: str,
|
||||
column: str,
|
||||
) -> bool:
|
||||
return any(
|
||||
str(row["name"]) == column
|
||||
for row in connection.execute(f"PRAGMA table_info({table})")
|
||||
)
|
||||
|
||||
|
||||
def migrate_database(connection: sqlite3.Connection) -> None:
|
||||
"""Bring a new or existing database to the latest known schema."""
|
||||
connection.execute("PRAGMA journal_mode=WAL")
|
||||
version = _schema_version(connection)
|
||||
|
||||
if version < 1:
|
||||
connection.executescript(INITIAL_SCHEMA)
|
||||
connection.execute("PRAGMA user_version = 1")
|
||||
version = 1
|
||||
|
||||
if version < 2:
|
||||
if not _has_column(connection, "user_settings", "yandex_model"):
|
||||
connection.execute(
|
||||
"ALTER TABLE user_settings ADD COLUMN yandex_model TEXT"
|
||||
)
|
||||
connection.execute("PRAGMA user_version = 2")
|
||||
@@ -3,9 +3,9 @@ import sqlite3
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Mapping
|
||||
|
||||
from .config import get_default_model
|
||||
from .migrations import migrate_database
|
||||
from .time_utils import to_utc_iso, utc_now
|
||||
|
||||
|
||||
@@ -21,8 +21,13 @@ def required_lastrowid(cursor: sqlite3.Cursor) -> int:
|
||||
|
||||
|
||||
class AssistantStorage:
|
||||
def __init__(self, path: Path) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
path: Path,
|
||||
default_models: Mapping[str, str],
|
||||
) -> None:
|
||||
self.path = path
|
||||
self._default_models = dict(default_models)
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._init_db()
|
||||
|
||||
@@ -45,92 +50,7 @@ class AssistantStorage:
|
||||
|
||||
def _init_db(self) -> None:
|
||||
with self._connection() as connection:
|
||||
connection.execute("PRAGMA journal_mode=WAL")
|
||||
connection.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS user_settings (
|
||||
user_id INTEGER PRIMARY KEY,
|
||||
ollama_model TEXT,
|
||||
yandex_model TEXT,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS authorized_users (
|
||||
user_id INTEGER PRIMARY KEY,
|
||||
authorized_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS memories (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS notes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS reminders (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
chat_id INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
remind_at TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
created_at TEXT NOT NULL,
|
||||
sent_at TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tracked_items (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversation_messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
chat_id INTEGER NOT NULL,
|
||||
role TEXT NOT NULL CHECK(role IN ('user', 'assistant')),
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversation_contexts (
|
||||
user_id INTEGER NOT NULL,
|
||||
chat_id INTEGER NOT NULL,
|
||||
started_after_id INTEGER NOT NULL DEFAULT 0,
|
||||
updated_at TEXT NOT NULL,
|
||||
PRIMARY KEY (user_id, chat_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_memories_user_created
|
||||
ON memories(user_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_notes_user_created
|
||||
ON notes(user_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_reminders_due
|
||||
ON reminders(status, remind_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_tracked_user_updated
|
||||
ON tracked_items(user_id, updated_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_chat_id
|
||||
ON conversation_messages(user_id, chat_id, id DESC);
|
||||
"""
|
||||
)
|
||||
columns = {
|
||||
str(row["name"])
|
||||
for row in connection.execute("PRAGMA table_info(user_settings)")
|
||||
}
|
||||
if "yandex_model" not in columns:
|
||||
connection.execute(
|
||||
"ALTER TABLE user_settings ADD COLUMN yandex_model TEXT"
|
||||
)
|
||||
migrate_database(connection)
|
||||
|
||||
def is_user_authorized(self, user_id: int) -> bool:
|
||||
with self._connection() as connection:
|
||||
@@ -312,7 +232,12 @@ class AssistantStorage:
|
||||
).fetchone()
|
||||
if row and row["model"]:
|
||||
return str(row["model"])
|
||||
return get_default_model(provider)
|
||||
try:
|
||||
return self._default_models[provider]
|
||||
except KeyError as exc:
|
||||
raise ValueError(
|
||||
f"No default model configured for provider: {provider}"
|
||||
) from exc
|
||||
|
||||
def set_user_model(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user