240 lines
10 KiB
Python
240 lines
10 KiB
Python
"""Database initialization: create all tables and seed defaults.
|
|
|
|
Idempotent: safe to call from multiple processes (backend lifespan + worker
|
|
startup) thanks to a PostgreSQL advisory lock that serializes the seeding
|
|
phase. `create_all` itself is already `CREATE TABLE IF NOT EXISTS`, so the
|
|
only race is on seed inserts — guarded by `pg_advisory_xact_lock` plus
|
|
per-row IntegrityError handling.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import select, text
|
|
from sqlalchemy.exc import IntegrityError
|
|
|
|
from app.db import AsyncSessionLocal, Base, engine
|
|
from app.models import GlossaryEntry, Preset, Setting, User
|
|
from app.config import settings
|
|
from app.logging_setup import get_logger, setup_logging
|
|
from app.core.security import hash_password
|
|
from app.prompts.fantasy_preset import FANTASY_PRESET_RU, FANTASY_PRESET_EN
|
|
|
|
log = get_logger("migrations")
|
|
|
|
|
|
# Stable advisory lock key so backend + worker don't race on seeding.
|
|
# (key1, key2) — arbitrary 64-bit integers, kept constant across runs.
|
|
_ADVISORY_LOCK_KEY = (42424201, 1)
|
|
|
|
|
|
DEFAULT_SETTINGS = [
|
|
("llm.base_url", settings.default_llm_base_url, "OpenAI-compatible base URL"),
|
|
("llm.api_key", settings.default_llm_api_key, "API key for LLM endpoint"),
|
|
("llm.model", settings.default_llm_model, "Default model name"),
|
|
("llm.temperature", 0.7, "Temperature for orchestrator"),
|
|
("llm.step_temperature", 0.85, "Temperature for narrative step writer"),
|
|
("llm.summary_temperature", 0.3, "Temperature for summarizer"),
|
|
("llm.max_tokens", 1024, "Max tokens per LLM response"),
|
|
("llm.request_timeout", 120, "LLM request timeout, seconds"),
|
|
("llm.streaming", True, "Whether to use streaming responses"),
|
|
("context.recent_messages", settings.default_recent_messages, "Guaranteed recent messages in prompt"),
|
|
("context.compress_threshold", settings.default_compress_threshold, "Trigger compression at this count"),
|
|
("context.summary_messages", settings.default_summary_messages, "Number of messages per summary block"),
|
|
("context.max_tokens_total", 6000, "Soft token budget for context window (small models)"),
|
|
("triggers.enabled", True, "Enable trigger firing on in-game time changes"),
|
|
# Embeddings / RAG
|
|
("embedding.provider", settings.default_embedding_provider, "Embeddings provider: 'hash' (offline fallback) or 'openai' (real semantic embeddings)"),
|
|
("embedding.base_url", settings.default_embedding_base_url, "OpenAI-compatible embeddings base URL. Empty = reuse llm.base_url"),
|
|
("embedding.api_key", settings.default_embedding_api_key, "API key for embeddings endpoint. Empty = reuse llm.api_key"),
|
|
("embedding.model", settings.default_embedding_model, "Embedding model name (e.g. text-embedding-3-small, bge-m3, nomic-embed-text)"),
|
|
("embedding.dim", settings.default_embedding_dim, "Vector dimension. 0 = auto-probe from endpoint on first use"),
|
|
("embedding.request_timeout", settings.default_embedding_request_timeout, "Embeddings request timeout, seconds"),
|
|
]
|
|
|
|
|
|
# Keys whose values come from environment variables (via Settings fields).
|
|
# These are re-applied on EVERY startup so .env is the source of truth.
|
|
# Admin-panel changes to these keys are runtime overrides that get reset on
|
|
# restart unless the operator also updates .env.
|
|
ENV_DERIVED_SETTING_KEYS = {
|
|
"llm.base_url",
|
|
"llm.api_key",
|
|
"llm.model",
|
|
"embedding.provider",
|
|
"embedding.base_url",
|
|
"embedding.api_key",
|
|
"embedding.model",
|
|
"embedding.dim",
|
|
"embedding.request_timeout",
|
|
}
|
|
|
|
|
|
async def init_db() -> None:
|
|
setup_logging()
|
|
log.info("creating_tables")
|
|
# CREATE TABLE IF NOT EXISTS — safe to run concurrently.
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
log.info("tables_ready")
|
|
|
|
# Seed phase: serialize across processes via PG advisory transaction lock.
|
|
# On non-PG backends (SQLite for tests) the lock statement is a no-op
|
|
# (we catch the error and proceed without locking).
|
|
try:
|
|
async with AsyncSessionLocal() as session:
|
|
await session.execute(
|
|
text("SELECT pg_advisory_xact_lock(:k1, :k2)").bindparams(
|
|
k1=_ADVISORY_LOCK_KEY[0], k2=_ADVISORY_LOCK_KEY[1]
|
|
)
|
|
)
|
|
await _seed_settings(session)
|
|
await _sync_env_derived_settings(session)
|
|
await _seed_builtin_presets(session)
|
|
await session.commit()
|
|
except Exception as e:
|
|
# Non-PG backend (SQLite) or transient error — retry without the lock.
|
|
log.warning("advisory_lock_unavailable_proceeding", error=f"{type(e).__name__}: {e}")
|
|
async with AsyncSessionLocal() as session:
|
|
await _seed_settings(session)
|
|
await _sync_env_derived_settings(session)
|
|
await _seed_builtin_presets(session)
|
|
await session.commit()
|
|
|
|
# Ensure admin_setup_token is set and print it on every startup.
|
|
#
|
|
# The admin-setup endpoint refuses to create a second admin (see app/api/auth.py),
|
|
# so it's safe to always print the token — even after an admin exists, the token
|
|
# is useless. We print on every startup (not just first run) so the operator can
|
|
# always find the URL in the logs without having to dig through old logs.
|
|
token = settings.admin_setup_token.strip()
|
|
if not token:
|
|
# No token forced via env — generate one and persist it (idempotent).
|
|
import secrets as _s
|
|
token = _s.token_urlsafe(24)
|
|
async with AsyncSessionLocal() as session:
|
|
existing = await session.execute(select(Setting).where(Setting.key == "admin.setup_token"))
|
|
existing_obj = existing.scalars().first()
|
|
if existing_obj is None:
|
|
session.add(Setting(key="admin.setup_token", value=token, description="One-time token for /admin/setup"))
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError:
|
|
# Another process inserted it concurrently — re-read.
|
|
await session.rollback()
|
|
await session.rollback()
|
|
existing = await session.execute(select(Setting).where(Setting.key == "admin.setup_token"))
|
|
existing_obj = existing.scalars().first()
|
|
if existing_obj is not None:
|
|
token = str(existing_obj.value)
|
|
else:
|
|
# Use the persisted token (env was empty, DB has one).
|
|
token = str(existing_obj.value)
|
|
# Always print — operator convenience.
|
|
print("=" * 60)
|
|
print("ADMIN SETUP URL:")
|
|
print(f" /admin/setup")
|
|
print("ADMIN SETUP TOKEN:")
|
|
print(f" {token}")
|
|
print("=" * 60)
|
|
log.info("admin_setup_token_printed")
|
|
|
|
|
|
async def _seed_settings(session) -> None:
|
|
"""Insert default settings that don't yet exist (per-row, race-safe)."""
|
|
result = await session.execute(select(Setting).limit(1))
|
|
if result.scalars().first() is not None:
|
|
log.info("settings_already_exist")
|
|
return
|
|
seeded = 0
|
|
for key, value, desc in DEFAULT_SETTINGS:
|
|
# Check existence per-row to avoid IntegrityError on concurrent inserts
|
|
existing = await session.execute(select(Setting).where(Setting.key == key))
|
|
if existing.scalars().first() is not None:
|
|
continue
|
|
session.add(Setting(key=key, value=value, description=desc))
|
|
seeded += 1
|
|
if seeded:
|
|
try:
|
|
await session.commit()
|
|
log.info("settings_seeded", count=seeded)
|
|
except IntegrityError:
|
|
await session.rollback()
|
|
log.info("settings_seed_skipped_concurrent")
|
|
else:
|
|
log.info("settings_already_exist")
|
|
|
|
|
|
async def _sync_env_derived_settings(session) -> None:
|
|
"""Re-apply env-derived setting values from .env on every startup.
|
|
|
|
This makes .env the source of truth for these keys: changing .env and
|
|
restarting the container takes effect immediately. Admin-panel edits to
|
|
these keys are runtime overrides that are reset on the next restart
|
|
(unless the operator also updates .env).
|
|
|
|
Only the env-derived keys (see ENV_DERIVED_SETTING_KEYS) are touched;
|
|
other settings (temperature, context params, etc.) are preserved as
|
|
configured via the admin panel.
|
|
"""
|
|
env_values = {key: value for key, value, _desc in DEFAULT_SETTINGS if key in ENV_DERIVED_SETTING_KEYS}
|
|
updated = 0
|
|
for key, new_value in env_values.items():
|
|
result = await session.execute(select(Setting).where(Setting.key == key))
|
|
row = result.scalars().first()
|
|
if row is None:
|
|
# Shouldn't happen (seeded above) but handle defensively.
|
|
session.add(Setting(key=key, value=new_value, description="Env-derived"))
|
|
updated += 1
|
|
else:
|
|
if row.value != new_value:
|
|
log.info(
|
|
"env_setting_resynced",
|
|
key=key,
|
|
old_value=str(row.value)[:80],
|
|
new_value=str(new_value)[:80],
|
|
)
|
|
row.value = new_value
|
|
updated += 1
|
|
if updated:
|
|
try:
|
|
await session.commit()
|
|
log.info("env_settings_synced", count=updated)
|
|
except IntegrityError:
|
|
await session.rollback()
|
|
log.warning("env_settings_sync_failed_concurrent")
|
|
|
|
|
|
async def _seed_builtin_presets(session) -> None:
|
|
"""Insert built-in presets if none exist yet."""
|
|
result = await session.execute(select(Preset).where(Preset.is_builtin.is_(True)))
|
|
if result.scalars().first() is not None:
|
|
log.info("builtin_presets_already_exist")
|
|
return
|
|
for preset_def in (FANTASY_PRESET_RU, FANTASY_PRESET_EN):
|
|
# Check by slug to avoid race on unique constraint
|
|
existing = await session.execute(select(Preset).where(Preset.slug == preset_def["slug"]))
|
|
if existing.scalars().first() is not None:
|
|
continue
|
|
session.add(Preset(
|
|
slug=preset_def["slug"],
|
|
title=preset_def["title"],
|
|
description=preset_def["description"],
|
|
language=preset_def["language"],
|
|
is_public=True,
|
|
is_builtin=True,
|
|
payload=preset_def["payload"],
|
|
))
|
|
try:
|
|
await session.commit()
|
|
log.info("builtin_presets_seeded")
|
|
except IntegrityError:
|
|
await session.rollback()
|
|
log.info("builtin_presets_seed_skipped_concurrent")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(init_db())
|