from collections.abc import AsyncGenerator import os import re from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase from app.config import settings os.makedirs(settings.DATA_DIR, exist_ok=True) engine = create_async_engine( settings.DATABASE_URL, echo=settings.DEBUG, pool_pre_ping=True, ) async_session = async_sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, ) class Base(DeclarativeBase): pass async def get_db() -> AsyncGenerator[AsyncSession, None]: async with async_session() as session: yield session async def init_db(): async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) await ensure_log_columns(conn) await ensure_message_columns(conn) await ensure_conversation_columns(conn) await ensure_document_columns(conn) await ensure_memory_columns(conn) await ensure_user_columns(conn) await ensure_forum_columns(conn) await ensure_agent_columns(conn) await ensure_skill_columns(conn) await ensure_learning_artifact_tables(conn) async def ensure_log_columns(conn): result = await conn.execute(text("PRAGMA table_info(logs)")) rows = result.fetchall() if not rows: return columns = {row[1] for row in rows} required_columns = { "request_id": "ALTER TABLE logs ADD COLUMN request_id VARCHAR(64)", "route": "ALTER TABLE logs ADD COLUMN route VARCHAR(255)", "method": "ALTER TABLE logs ADD COLUMN method VARCHAR(16)", "status_code": "ALTER TABLE logs ADD COLUMN status_code INTEGER", "error_type": "ALTER TABLE logs ADD COLUMN error_type VARCHAR(100)", "operation": "ALTER TABLE logs ADD COLUMN operation VARCHAR(100)", } for column, ddl in required_columns.items(): if column not in columns: await conn.execute(text(ddl)) async def ensure_message_columns(conn): result = await conn.execute(text("PRAGMA table_info(messages)")) rows = result.fetchall() if not rows: return columns = {row[1] for row in rows} required_columns = { "attachments": "ALTER TABLE messages ADD COLUMN attachments JSON", } for column, ddl in required_columns.items(): if column not in columns: await conn.execute(text(ddl)) async def ensure_conversation_columns(conn): rows = await _get_table_info(conn, 'conversations') if not rows: return columns = {row[1] for row in rows} required_columns = { 'agent_state': "ALTER TABLE conversations ADD COLUMN agent_state JSON", } for column, ddl in required_columns.items(): if column not in columns: await conn.execute(text(ddl)) async def ensure_document_columns(conn): result = await conn.execute(text("PRAGMA table_info(documents)")) rows = result.fetchall() if not rows: return columns = {row[1] for row in rows} required_columns = { "ingestion_status": "ALTER TABLE documents ADD COLUMN ingestion_status VARCHAR(50) DEFAULT 'uploaded' NOT NULL", "ingestion_error": "ALTER TABLE documents ADD COLUMN ingestion_error TEXT", "indexed_at": "ALTER TABLE documents ADD COLUMN indexed_at DATETIME", "parser_version": "ALTER TABLE documents ADD COLUMN parser_version VARCHAR(50)", "index_version": "ALTER TABLE documents ADD COLUMN index_version VARCHAR(50)", "normalized_content": "ALTER TABLE documents ADD COLUMN normalized_content TEXT", "normalized_format": "ALTER TABLE documents ADD COLUMN normalized_format VARCHAR(50)", } for column, ddl in required_columns.items(): if column not in columns: await conn.execute(text(ddl)) async def ensure_memory_columns(conn): rows = await _get_table_info(conn, 'user_memories') if not rows: return columns = {row[1] for row in rows} required_columns = { 'frequency_count': "ALTER TABLE user_memories ADD COLUMN frequency_count INTEGER DEFAULT 0", 'emotion_tags': "ALTER TABLE user_memories ADD COLUMN emotion_tags JSON", 'importance_score': "ALTER TABLE user_memories ADD COLUMN importance_score FLOAT DEFAULT 0.5", 'importance_level': "ALTER TABLE user_memories ADD COLUMN importance_level VARCHAR(20) DEFAULT 'medium'", 'associated_topics': "ALTER TABLE user_memories ADD COLUMN associated_topics JSON", 'decay_score': "ALTER TABLE user_memories ADD COLUMN decay_score FLOAT DEFAULT 1.0", 'is_archived': "ALTER TABLE user_memories ADD COLUMN is_archived BOOLEAN DEFAULT 0", 'last_accessed_at': "ALTER TABLE user_memories ADD COLUMN last_accessed_at DATETIME", 'archive_at': "ALTER TABLE user_memories ADD COLUMN archive_at DATETIME", } for column, ddl in required_columns.items(): if column not in columns: await conn.execute(text(ddl)) async def ensure_user_columns(conn): rows = await _get_table_info(conn, 'users') if not rows: return columns = {row[1] for row in rows} if 'username' not in columns: await conn.execute(text("ALTER TABLE users ADD COLUMN username VARCHAR(255)")) rows = await _get_table_info(conn, 'users') await _backfill_usernames(conn) username_row = next(row for row in rows if row[1] == 'username') indexes = await _get_index_info(conn, 'users') has_username_index = any(row[1] == 'ix_users_username' and row[2] == 1 for row in indexes) has_email_index = any(row[1] == 'ix_users_email' and row[2] == 1 for row in indexes) if username_row[3] != 1 or not has_username_index or not has_email_index: await _rebuild_users_table(conn) async def ensure_forum_columns(conn): rows = await _get_table_info(conn, 'forum_posts') if not rows: return columns = {row[1] for row in rows} required_columns = { "board": "ALTER TABLE forum_posts ADD COLUMN board VARCHAR(100) DEFAULT 'general' NOT NULL", "is_pinned": "ALTER TABLE forum_posts ADD COLUMN is_pinned BOOLEAN DEFAULT 0 NOT NULL", } for column, ddl in required_columns.items(): if column not in columns: await conn.execute(text(ddl)) indexes = await _get_index_info(conn, 'forum_posts') index_names = {row[1] for row in indexes} if 'ix_forum_posts_board' not in index_names: await conn.execute(text("CREATE INDEX IF NOT EXISTS ix_forum_posts_board ON forum_posts (board)")) async def ensure_agent_columns(conn): rows = await _get_table_info(conn, 'agents') if not rows: return columns = {row[1] for row in rows} required_columns = { 'selected_skill_ids': "ALTER TABLE agents ADD COLUMN selected_skill_ids JSON DEFAULT '[]' NOT NULL", } for column, ddl in required_columns.items(): if column not in columns: await conn.execute(text(ddl)) async def ensure_skill_columns(conn): rows = await _get_table_info(conn, 'skills') if not rows: return columns = {row[1] for row in rows} required_columns = { 'required_context': "ALTER TABLE skills ADD COLUMN required_context JSON DEFAULT '[]' NOT NULL", 'output_format': "ALTER TABLE skills ADD COLUMN output_format TEXT", 'is_builtin': "ALTER TABLE skills ADD COLUMN is_builtin BOOLEAN DEFAULT 0 NOT NULL", 'team_id': "ALTER TABLE skills ADD COLUMN team_id VARCHAR(36)", 'status': "ALTER TABLE skills ADD COLUMN status VARCHAR(20) DEFAULT 'active' NOT NULL", 'scope': "ALTER TABLE skills ADD COLUMN scope JSON DEFAULT '[]' NOT NULL", 'effectiveness': "ALTER TABLE skills ADD COLUMN effectiveness FLOAT DEFAULT 0.0 NOT NULL", 'review_after': "ALTER TABLE skills ADD COLUMN review_after DATETIME", 'candidate_count': "ALTER TABLE skills ADD COLUMN candidate_count INTEGER DEFAULT 0 NOT NULL", 'candidate_source_hashes': "ALTER TABLE skills ADD COLUMN candidate_source_hashes JSON DEFAULT '[]' NOT NULL", 'activation_count': "ALTER TABLE skills ADD COLUMN activation_count INTEGER DEFAULT 0 NOT NULL", 'last_activated_at': "ALTER TABLE skills ADD COLUMN last_activated_at DATETIME", } for column, ddl in required_columns.items(): if column not in columns: await conn.execute(text(ddl)) await conn.execute(text("UPDATE skills SET agent_type = 'schedule_planner' WHERE agent_type = 'planner'")) builtin_names = [ '今日重点拆解', '周计划编排', '时间冲突分析', '任务执行 SOP', '外部交互推进', '知识检索摘要', '图谱沉淀策略', '风险识别模板', '趋势洞察模板', ] for name in builtin_names: await conn.execute( text("UPDATE skills SET is_builtin = 1 WHERE name = :name"), {'name': name}, ) async def ensure_learning_artifact_tables(conn): await conn.execute( text( """ CREATE TABLE IF NOT EXISTS learning_artifacts ( id VARCHAR(36) PRIMARY KEY, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, user_id VARCHAR(36) NOT NULL, conversation_id VARCHAR(36) NOT NULL, retrospective_id VARCHAR(36), artifact_type VARCHAR(32) NOT NULL, artifact_key VARCHAR(128), summary_text TEXT NOT NULL, payload JSON NOT NULL, recorded_at DATETIME NOT NULL ) """ ) ) await conn.execute( text( "CREATE INDEX IF NOT EXISTS ix_learning_artifacts_user_id ON learning_artifacts (user_id)" ) ) await conn.execute( text( "CREATE INDEX IF NOT EXISTS ix_learning_artifacts_conversation_id ON learning_artifacts (conversation_id)" ) ) await conn.execute( text( "CREATE INDEX IF NOT EXISTS ix_learning_artifacts_retrospective_id ON learning_artifacts (retrospective_id)" ) ) await conn.execute( text( "CREATE INDEX IF NOT EXISTS ix_learning_artifacts_artifact_type ON learning_artifacts (artifact_type)" ) ) async def _backfill_usernames(conn): result = await conn.execute(text("SELECT id, email, username FROM users ORDER BY created_at, id")) users = result.fetchall() seen_usernames: set[str] = set() for user_id, email, username in users: if username: seen_usernames.add(username) continue base_username = _slugify_username((email or '').split('@', 1)[0]) candidate = base_username suffix = 2 while candidate in seen_usernames: candidate = f"{base_username}_{suffix}" suffix += 1 await conn.execute( text("UPDATE users SET username = :username WHERE id = :user_id AND username IS NULL"), {"username": candidate, "user_id": user_id}, ) seen_usernames.add(candidate) async def _rebuild_users_table(conn): await conn.execute(text("CREATE TABLE users__new (id VARCHAR(36) PRIMARY KEY, username VARCHAR(255) NOT NULL, email VARCHAR(255) NOT NULL, hashed_password VARCHAR(255) NOT NULL, full_name VARCHAR(255), is_active BOOLEAN NOT NULL DEFAULT 1, is_superuser BOOLEAN NOT NULL DEFAULT 0, llm_config JSON, scheduler_config JSON, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)")) await conn.execute(text("INSERT INTO users__new (id, username, email, hashed_password, full_name, is_active, is_superuser, llm_config, scheduler_config, created_at, updated_at) SELECT id, username, email, hashed_password, full_name, COALESCE(is_active, 1), COALESCE(is_superuser, 0), llm_config, scheduler_config, COALESCE(created_at, CURRENT_TIMESTAMP), COALESCE(updated_at, CURRENT_TIMESTAMP) FROM users")) await conn.execute(text("DROP TABLE users")) await conn.execute(text("ALTER TABLE users__new RENAME TO users")) await conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS ix_users_username ON users (username)")) await conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS ix_users_email ON users (email)")) async def _get_table_info(conn, table_name: str): result = await conn.execute(text(f"PRAGMA table_info({table_name})")) return result.fetchall() async def _get_index_info(conn, table_name: str): result = await conn.execute(text(f"PRAGMA index_list({table_name})")) return result.fetchall() def _slugify_username(value: str) -> str: normalized = re.sub(r'[^a-z0-9_]+', '_', value.strip().lower()) normalized = re.sub(r'_+', '_', normalized).strip('_') return normalized or 'user'