101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
from sqlalchemy import text
|
|
from sqlmodel import SQLModel, Session, create_engine
|
|
|
|
from core.settings import DATABASE_ECHO, DATABASE_URL
|
|
|
|
# Ensure table models are registered in SQLModel metadata before create_all.
|
|
from models import bot as _bot_models # noqa: F401
|
|
|
|
engine = create_engine(DATABASE_URL, echo=DATABASE_ECHO)
|
|
|
|
|
|
def _ensure_botinstance_columns() -> None:
|
|
if engine.dialect.name != "sqlite":
|
|
return
|
|
required_columns = {
|
|
"current_state": "TEXT DEFAULT 'IDLE'",
|
|
"last_action": "TEXT",
|
|
"image_tag": "TEXT DEFAULT 'nanobot-base:v0.1.4'",
|
|
}
|
|
with engine.connect() as conn:
|
|
existing_rows = conn.execute(text("PRAGMA table_info(botinstance)")).fetchall()
|
|
existing = {str(row[1]) for row in existing_rows}
|
|
for col, ddl in required_columns.items():
|
|
if col in existing:
|
|
continue
|
|
conn.execute(text(f"ALTER TABLE botinstance ADD COLUMN {col} {ddl}"))
|
|
conn.commit()
|
|
|
|
|
|
def _drop_legacy_botinstance_columns() -> None:
|
|
if engine.dialect.name != "sqlite":
|
|
return
|
|
legacy_columns = [
|
|
"avatar_model",
|
|
"avatar_skin",
|
|
"system_prompt",
|
|
"soul_md",
|
|
"agents_md",
|
|
"user_md",
|
|
"tools_md",
|
|
"tools_config_json",
|
|
"identity_md",
|
|
"llm_provider",
|
|
"llm_model",
|
|
"api_key",
|
|
"api_base",
|
|
"temperature",
|
|
"top_p",
|
|
"max_tokens",
|
|
"presence_penalty",
|
|
"frequency_penalty",
|
|
"send_progress",
|
|
"send_tool_hints",
|
|
"bot_env_json",
|
|
]
|
|
with engine.connect() as conn:
|
|
existing_rows = conn.execute(text("PRAGMA table_info(botinstance)")).fetchall()
|
|
existing = {str(row[1]) for row in existing_rows}
|
|
for col in legacy_columns:
|
|
if col not in existing:
|
|
continue
|
|
conn.execute(text(f'ALTER TABLE botinstance DROP COLUMN "{col}"'))
|
|
conn.commit()
|
|
|
|
|
|
def _ensure_botmessage_columns() -> None:
|
|
if engine.dialect.name != "sqlite":
|
|
return
|
|
required_columns = {
|
|
"media_json": "TEXT",
|
|
}
|
|
with engine.connect() as conn:
|
|
existing_rows = conn.execute(text("PRAGMA table_info(botmessage)")).fetchall()
|
|
existing = {str(row[1]) for row in existing_rows}
|
|
for col, ddl in required_columns.items():
|
|
if col in existing:
|
|
continue
|
|
conn.execute(text(f"ALTER TABLE botmessage ADD COLUMN {col} {ddl}"))
|
|
conn.commit()
|
|
|
|
|
|
def _drop_legacy_skill_tables() -> None:
|
|
"""Drop deprecated skill registry tables (moved to workspace filesystem mode)."""
|
|
with engine.connect() as conn:
|
|
conn.execute(text("DROP TABLE IF EXISTS botskillmapping"))
|
|
conn.execute(text("DROP TABLE IF EXISTS skillregistry"))
|
|
conn.commit()
|
|
|
|
|
|
def init_database() -> None:
|
|
SQLModel.metadata.create_all(engine)
|
|
_drop_legacy_skill_tables()
|
|
_ensure_botinstance_columns()
|
|
_drop_legacy_botinstance_columns()
|
|
_ensure_botmessage_columns()
|
|
|
|
|
|
def get_session():
|
|
with Session(engine) as session:
|
|
yield session
|