"""创建 Agent 库基础表。"""

from alembic import op

from zbt.infrastructure.mysql.agent_models import AgentBase

revision = "agent_0001"
down_revision = None
branch_labels = None
depends_on = None

TABLE_NAMES = ("agent_threads", "agent_messages", "agent_runs")


def upgrade() -> None:
    bind = op.get_bind()
    for table in AgentBase.metadata.sorted_tables:
        if table.name in TABLE_NAMES:
            table.create(bind=bind)


def downgrade() -> None:
    bind = op.get_bind()
    for table in reversed(AgentBase.metadata.sorted_tables):
        if table.name in TABLE_NAMES:
            table.drop(bind=bind)