added RAG, Multiuser, TG bot
This commit is contained in:
@@ -0,0 +1,249 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import secrets
|
||||
from pathlib import Path
|
||||
|
||||
from sqlalchemy import inspect, text
|
||||
|
||||
from app.auth.tokens import hash_token
|
||||
from app.character.card import DEFAULT_CARD, normalize_card
|
||||
from app.config import get_settings
|
||||
from app.db.base import engine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TENANT_TABLES = (
|
||||
"chat_sessions",
|
||||
"user_profile",
|
||||
"memory_facts",
|
||||
"fitness_profiles",
|
||||
"body_metrics",
|
||||
"food_logs",
|
||||
"water_logs",
|
||||
"workout_logs",
|
||||
"step_logs",
|
||||
"fitness_reminders",
|
||||
"shopping_lists",
|
||||
"reminders",
|
||||
"documents",
|
||||
"pomodoro_cycles",
|
||||
"pomodoro_sessions",
|
||||
"project_bindings",
|
||||
"work_items",
|
||||
)
|
||||
|
||||
LEGACY_CARD_PATH = Path("./data/character.json")
|
||||
|
||||
|
||||
def _table_exists(name: str) -> bool:
|
||||
return name in inspect(engine).get_table_names()
|
||||
|
||||
|
||||
def _columns(table: str) -> set[str]:
|
||||
if not _table_exists(table):
|
||||
return set()
|
||||
return {col["name"] for col in inspect(engine).get_columns(table)}
|
||||
|
||||
|
||||
def _add_column_if_missing(table: str, column: str, ddl: str) -> None:
|
||||
if column in _columns(table):
|
||||
return
|
||||
with engine.begin() as conn:
|
||||
conn.execute(text(ddl))
|
||||
|
||||
|
||||
def _ensure_users_table() -> None:
|
||||
if _table_exists("users"):
|
||||
return
|
||||
with engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(
|
||||
"CREATE TABLE users ("
|
||||
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
|
||||
"username VARCHAR(64) NOT NULL UNIQUE, "
|
||||
"display_name VARCHAR(255) DEFAULT '', "
|
||||
"api_token_hash VARCHAR(64) NOT NULL, "
|
||||
"is_active BOOLEAN DEFAULT 1, "
|
||||
"created_at DATETIME DEFAULT CURRENT_TIMESTAMP"
|
||||
")"
|
||||
)
|
||||
)
|
||||
conn.execute(text("CREATE INDEX IF NOT EXISTS ix_users_api_token_hash ON users (api_token_hash)"))
|
||||
|
||||
|
||||
def _ensure_character_cards_table() -> None:
|
||||
if _table_exists("character_cards"):
|
||||
return
|
||||
with engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(
|
||||
"CREATE TABLE character_cards ("
|
||||
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
|
||||
"user_id INTEGER NOT NULL UNIQUE REFERENCES users(id) ON DELETE CASCADE, "
|
||||
"card_json TEXT DEFAULT '{}', "
|
||||
"updated_at DATETIME DEFAULT CURRENT_TIMESTAMP"
|
||||
")"
|
||||
)
|
||||
)
|
||||
conn.execute(text("CREATE INDEX IF NOT EXISTS ix_character_cards_user_id ON character_cards (user_id)"))
|
||||
|
||||
|
||||
def _add_user_id_columns() -> None:
|
||||
for table in TENANT_TABLES:
|
||||
if not _table_exists(table):
|
||||
continue
|
||||
_add_column_if_missing(
|
||||
table,
|
||||
"user_id",
|
||||
f"ALTER TABLE {table} ADD COLUMN user_id INTEGER REFERENCES users(id) ON DELETE CASCADE",
|
||||
)
|
||||
with engine.begin() as conn:
|
||||
conn.execute(text(f"CREATE INDEX IF NOT EXISTS ix_{table}_user_id ON {table} (user_id)"))
|
||||
|
||||
|
||||
def _ensure_default_user() -> tuple[int, str | None]:
|
||||
settings = get_settings()
|
||||
with engine.begin() as conn:
|
||||
row = conn.execute(text("SELECT id FROM users ORDER BY id LIMIT 1")).fetchone()
|
||||
if row:
|
||||
return int(row[0]), None
|
||||
|
||||
username = settings.default_user_username or "owner"
|
||||
display_name = settings.default_user_display_name or username
|
||||
plain_token = (settings.default_api_token or "").strip()
|
||||
generated = False
|
||||
if not plain_token:
|
||||
plain_token = secrets.token_urlsafe(32)
|
||||
generated = True
|
||||
|
||||
token_hash = hash_token(plain_token)
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO users (id, username, display_name, api_token_hash, is_active) "
|
||||
"VALUES (1, :username, :display_name, :token_hash, 1)"
|
||||
),
|
||||
{"username": username, "display_name": display_name, "token_hash": token_hash},
|
||||
)
|
||||
if generated:
|
||||
logger.warning(
|
||||
"DEFAULT_API_TOKEN not set — generated token for user '%s': %s",
|
||||
username,
|
||||
plain_token,
|
||||
)
|
||||
return 1, plain_token
|
||||
return 1, None
|
||||
|
||||
|
||||
def _backfill_user_id(default_user_id: int = 1) -> None:
|
||||
with engine.begin() as conn:
|
||||
for table in TENANT_TABLES:
|
||||
if not _table_exists(table):
|
||||
continue
|
||||
conn.execute(
|
||||
text(f"UPDATE {table} SET user_id = :uid WHERE user_id IS NULL"),
|
||||
{"uid": default_user_id},
|
||||
)
|
||||
|
||||
|
||||
def _rebuild_shopping_unique() -> None:
|
||||
if not _table_exists("shopping_lists"):
|
||||
return
|
||||
with engine.begin() as conn:
|
||||
conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS uq_shopping_lists_user_name ON shopping_lists (user_id, name)"))
|
||||
|
||||
|
||||
def _rebuild_project_bindings_unique() -> None:
|
||||
if not _table_exists("project_bindings"):
|
||||
return
|
||||
with engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS uq_project_bindings_user_slug "
|
||||
"ON project_bindings (user_id, taiga_slug)"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _import_character_card(user_id: int) -> None:
|
||||
with engine.begin() as conn:
|
||||
existing = conn.execute(
|
||||
text("SELECT id FROM character_cards WHERE user_id = :uid"),
|
||||
{"uid": user_id},
|
||||
).fetchone()
|
||||
if existing:
|
||||
return
|
||||
|
||||
card = normalize_card(DEFAULT_CARD)
|
||||
if LEGACY_CARD_PATH.is_file():
|
||||
try:
|
||||
raw = json.loads(LEGACY_CARD_PATH.read_text(encoding="utf-8"))
|
||||
card = normalize_card(raw)
|
||||
except (json.JSONDecodeError, OSError):
|
||||
pass
|
||||
|
||||
conn.execute(
|
||||
text("INSERT INTO character_cards (user_id, card_json) VALUES (:uid, :json)"),
|
||||
{"uid": user_id, "json": json.dumps(card, ensure_ascii=False)},
|
||||
)
|
||||
|
||||
|
||||
def _backfill_qdrant_user_id(default_user_id: int = 1) -> None:
|
||||
settings = get_settings()
|
||||
if not settings.rag_enabled:
|
||||
return
|
||||
try:
|
||||
from app.rag.store import COLLECTION_DOC_CHUNKS, COLLECTION_FACTS, COLLECTION_SUMMARIES, _client
|
||||
except Exception:
|
||||
logger.exception("Qdrant backfill skipped")
|
||||
return
|
||||
|
||||
try:
|
||||
client = _client()
|
||||
except Exception:
|
||||
logger.warning('Qdrant unavailable, skipping user_id backfill')
|
||||
return
|
||||
|
||||
for collection in (COLLECTION_FACTS, COLLECTION_SUMMARIES, COLLECTION_DOC_CHUNKS):
|
||||
try:
|
||||
if not client.collection_exists(collection):
|
||||
continue
|
||||
except Exception:
|
||||
logger.warning('Qdrant unavailable for collection %s', collection)
|
||||
continue
|
||||
offset = None
|
||||
while True:
|
||||
points, offset = client.scroll(
|
||||
collection_name=collection,
|
||||
limit=100,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
with_vectors=False,
|
||||
)
|
||||
if not points:
|
||||
break
|
||||
missing = [point.id for point in points if (point.payload or {}).get("user_id") is None]
|
||||
if missing:
|
||||
client.set_payload(
|
||||
collection_name=collection,
|
||||
payload={"user_id": default_user_id},
|
||||
points=missing,
|
||||
)
|
||||
if offset is None:
|
||||
break
|
||||
logger.info("Qdrant user_id backfill completed for user_id=%s", default_user_id)
|
||||
|
||||
|
||||
def run_multi_user_migrations() -> str | None:
|
||||
"""Returns newly generated API token if any."""
|
||||
_ensure_users_table()
|
||||
_ensure_character_cards_table()
|
||||
_add_user_id_columns()
|
||||
user_id, new_token = _ensure_default_user()
|
||||
_backfill_user_id(user_id)
|
||||
_rebuild_shopping_unique()
|
||||
_rebuild_project_bindings_unique()
|
||||
_import_character_card(user_id)
|
||||
_backfill_qdrant_user_id(user_id)
|
||||
return new_token
|
||||
Reference in New Issue
Block a user