"""AI Labs — conversations / messages / generations / usage + 2 cron seeds. Four tables. The split between the first three (employee content, hard deleted on request) and the fourth (counters only, what admin surfaces read) is the domain's privacy design, documented on the models. ``ai_labs_messages`` is the one table here with a real ``ON DELETE CASCADE``: deleting a thread must take its turns with it even if a future code path forgets, because "Helix still has the text after telling you it did not" is the failure this domain cannot have. The service also deletes the children explicitly, since TiDB deployments vary in FK enforcement — belt and braces on the one thing that must not be half-done. Revision ID: 601 Revises: 600 """ from __future__ import annotations import uuid import sqlalchemy as sa from alembic import op from app.core.migration_helpers import id_type_for, table_kw_for revision = "601" down_revision = "600" branch_labels = None depends_on = None T_CONV = "ai_labs_conversations" T_MSG = "ai_labs_messages" T_GEN = "ai_labs_generations" T_USAGE = "ai_labs_usage_daily" JOB_REAP = "ai_labs.reap_generations" JOB_RETENTION = "ai_labs.sweep_retention" def _table_exists(bind, name: str) -> bool: return bind.execute( sa.text( "SELECT 1 FROM information_schema.tables " "WHERE table_schema = DATABASE() AND table_name = :n" ), {"n": name}, ).scalar() is not None def _seed_job( bind, *, name: str, display_name: str, description: str, cron: str, ) -> None: exists = bind.execute( sa.text("SELECT 1 FROM scheduled_jobs WHERE name = :n LIMIT 1"), {"n": name}, ).scalar() if exists: return bind.execute( sa.text( """ INSERT INTO scheduled_jobs (id, name, display_name, description, handler_key, trigger_type, interval_seconds, cron_expr, timezone, enabled, config, run_count, failure_count, created_at, updated_at) VALUES (:id, :name, :display_name, :description, :name, 'cron', NULL, :cron, 'Asia/Shanghai', 1, '{}', 0, 0, NOW(), NOW()) """ ), { "id": str(uuid.uuid4()), "name": name, "display_name": display_name, "description": description, "cron": cron, }, ) def upgrade() -> None: bind = op.get_bind() # ``id_type_for`` pins the parent's COLLATION as well as its type: # MySQL/TiDB raise errno 3780 on a collation mismatch alone, and that # only reproduces on TiDB. pid = id_type_for("identity_persons") kw = table_kw_for("identity_persons") if not _table_exists(bind, T_CONV): op.create_table( T_CONV, sa.Column("id", pid, primary_key=True), sa.Column( "person_id", pid, sa.ForeignKey("identity_persons.id"), nullable=False, ), sa.Column("title", sa.String(200), nullable=False, server_default=""), sa.Column("model_key", sa.String(64), nullable=False), sa.Column( "thinking_enabled", sa.Boolean(), nullable=False, server_default=sa.text("0"), ), sa.Column( "pinned", sa.Boolean(), nullable=False, server_default=sa.text("0"), ), sa.Column( "message_count", sa.Integer(), nullable=False, server_default=sa.text("0"), ), sa.Column("last_message_at", sa.DateTime(timezone=True)), sa.Column( "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now(), ), sa.Column( "updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now(), ), **kw, ) op.create_index( "ix_alc_person_activity", T_CONV, ["person_id", "last_message_at"], ) op.create_index( "ix_alc_person_created", T_CONV, ["person_id", "created_at"], ) if not _table_exists(bind, T_MSG): op.create_table( T_MSG, sa.Column("id", pid, primary_key=True), sa.Column( "conversation_id", pid, sa.ForeignKey(f"{T_CONV}.id", ondelete="CASCADE"), nullable=False, ), sa.Column("seq", sa.Integer(), nullable=False), sa.Column("role", sa.String(16), nullable=False), sa.Column("content", sa.Text(), nullable=False), sa.Column("reasoning", sa.Text()), sa.Column("attachments", sa.JSON()), sa.Column("model", sa.String(64)), sa.Column("input_tokens", sa.Integer()), sa.Column("output_tokens", sa.Integer()), sa.Column("latency_ms", sa.Integer()), sa.Column("finish_reason", sa.String(32)), sa.Column("error_code", sa.String(64)), sa.Column("client_token", sa.String(128)), sa.Column( "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now(), ), sa.UniqueConstraint( "conversation_id", "seq", name="uq_alm_conv_seq", ), # What stops a RETRY of one send from becoming a second turn # and a second paid vendor call. ``uq_alm_conv_seq`` only # catches the concurrent case: a sequential retry reads a # fresh ``message_count`` and would be accepted. Nullable, so # a caller that sends no token behaves exactly as before. sa.UniqueConstraint( "conversation_id", "client_token", name="uq_alm_conv_client", ), **kw, ) if not _table_exists(bind, T_GEN): op.create_table( T_GEN, sa.Column("id", pid, primary_key=True), sa.Column( "person_id", pid, sa.ForeignKey("identity_persons.id"), nullable=False, ), sa.Column("capability", sa.String(32), nullable=False), sa.Column("model", sa.String(64), nullable=False), sa.Column("status", sa.String(16), nullable=False), sa.Column("prompt", sa.Text(), nullable=False), sa.Column("params", sa.JSON()), sa.Column("input_object_keys", sa.JSON()), sa.Column("artifacts", sa.JSON()), sa.Column("output_text", sa.Text()), sa.Column("vendor_task_id", sa.String(128)), sa.Column("error_code", sa.String(64)), sa.Column("error_message", sa.String(1000)), sa.Column( "attempts", sa.Integer(), nullable=False, server_default=sa.text("0"), ), sa.Column("lease_owner", sa.String(64)), sa.Column("lease_expires_at", sa.DateTime(timezone=True)), sa.Column("started_at", sa.DateTime(timezone=True)), sa.Column("finished_at", sa.DateTime(timezone=True)), sa.Column("last_polled_at", sa.DateTime(timezone=True)), sa.Column( "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now(), ), sa.Column( "updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now(), ), **kw, ) op.create_index( "ix_alg_person_created", T_GEN, ["person_id", "created_at"], ) op.create_index( "ix_alg_person_capability", T_GEN, ["person_id", "capability"], ) op.create_index( "ix_alg_status_created", T_GEN, ["status", "created_at"], ) op.create_index("ix_alg_vendor_task", T_GEN, ["vendor_task_id"]) if not _table_exists(bind, T_USAGE): # No FK on ``person_id`` here on purpose: these counters must # outlive the content they describe (and a hard-deleted person's # historical usage is still part of the platform's own totals). op.create_table( T_USAGE, sa.Column("id", pid, primary_key=True), sa.Column("person_id", pid, nullable=False), sa.Column("usage_date", sa.Date(), nullable=False), sa.Column("capability", sa.String(32), nullable=False), sa.Column("model", sa.String(64), nullable=False), sa.Column( "requests", sa.Integer(), nullable=False, server_default=sa.text("0"), ), sa.Column( "failures", sa.Integer(), nullable=False, server_default=sa.text("0"), ), sa.Column( "input_tokens", sa.Integer(), nullable=False, server_default=sa.text("0"), ), sa.Column( "output_tokens", sa.Integer(), nullable=False, server_default=sa.text("0"), ), sa.Column( "artifacts", sa.Integer(), nullable=False, server_default=sa.text("0"), ), sa.Column( "media_seconds", sa.Integer(), nullable=False, server_default=sa.text("0"), ), sa.Column( "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now(), ), sa.Column( "updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now(), ), sa.UniqueConstraint( "person_id", "usage_date", "capability", "model", name="uq_alud_person_date_cap_model", ), **kw, ) op.create_index("ix_alud_date", T_USAGE, ["usage_date"]) op.create_index( "ix_alud_person_date", T_USAGE, ["person_id", "usage_date"], ) _seed_job( bind, name=JOB_REAP, display_name="AI 实验室 生成任务续跑", description=( "每 5 分钟找出租约已过期但未完成的生成任务并继续驱动;" "重试次数用尽的判失败。" ), # Every five minutes: a job orphaned by a deploy should resume # within one coffee sip, not one cron hour. cron="*/5 * * * *", ) _seed_job( bind, name=JOB_RETENTION, display_name="AI 实验室 历史保留期清理", description=( "当 ai_labs.retention_days > 0 时清理超期的对话与作品(含 S3 " "文件);默认 0 = 不清理。" ), cron="40 3 * * *", ) def downgrade() -> None: bind = op.get_bind() bind.execute( sa.text("DELETE FROM scheduled_jobs WHERE name IN (:a, :b)"), {"a": JOB_REAP, "b": JOB_RETENTION}, ) # Children before parents: ``ai_labs_messages`` references # ``ai_labs_conversations``. for table in (T_MSG, T_CONV, T_GEN, T_USAGE): if _table_exists(bind, table): op.drop_table(table)