SupplyAgent/alembic/versions/20260727_03_agent_document_injection.py

"""create persisted Agent document injection table

Revision ID: 20260727_03
Revises: 20260727_02
Create Date: 2026-07-27
"""
from __future__ import annotations

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import mysql

revision: str = "20260727_03"
down_revision: str | None = "20260727_02"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    op.create_table(
        "agent_document_injection",
        sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column(
            "agent_name",
            sa.String(length=128),
            nullable=False,
            comment="Agent name",
        ),
        sa.Column(
            "content",
            sa.Text().with_variant(mysql.LONGTEXT(), "mysql"),
            nullable=False,
            comment="Document content injected after the base system prompt",
        ),
        sa.Column(
            "enabled",
            sa.Integer(),
            server_default=sa.text("1"),
            nullable=False,
            comment="Whether the document injection is active",
        ),
        sa.Column(
            "create_time",
            sa.DateTime(),
            server_default=sa.func.now(),
            nullable=False,
        ),
        sa.Column(
            "update_time",
            sa.DateTime(),
            server_default=sa.func.now(),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "agent_name",
            name="uk_agent_document_injection_agent_name",
        ),
        mysql_charset="utf8mb4",
    )
    op.create_index(
        "ix_agent_document_injection_agent_name",
        "agent_document_injection",
        ["agent_name"],
    )


def downgrade() -> None:
    op.drop_index(
        "ix_agent_document_injection_agent_name",
        table_name="agent_document_injection",
    )
    op.drop_table("agent_document_injection")