openassistant/migrations/versions/2024-04-22-17-19_aa4bda3363e3.py

"""update models

Revision ID: aa4bda3363e3
Revises: 8dbb8f38ef77
Create Date: 2024-04-22 17:19:59.829072

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import sqlmodel


# revision identifiers, used by Alembic.
revision: str = "aa4bda3363e3"
down_revision: Union[str, None] = "8dbb8f38ef77"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("assistant", sa.Column("response_format", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
    op.add_column("assistant", sa.Column("tool_resources", sa.JSON(), nullable=True))
    op.add_column("assistant", sa.Column("temperature", sa.Float(), nullable=True))
    op.add_column("assistant", sa.Column("top_p", sa.Float(), nullable=True))
    op.add_column("run", sa.Column("incomplete_details", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
    op.add_column("run", sa.Column("max_completion_tokens", sa.Integer(), nullable=True))
    op.add_column("run", sa.Column("max_prompt_tokens", sa.Integer(), nullable=True))
    op.add_column("run", sa.Column("response_format", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
    op.add_column("run", sa.Column("tool_choice", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
    op.add_column("run", sa.Column("truncation_strategy", sa.JSON(), nullable=True))
    op.add_column("run", sa.Column("usage", sa.JSON(), nullable=True))
    op.add_column("run", sa.Column("temperature", sa.Float(), nullable=True))
    op.add_column("run", sa.Column("top_p", sa.Float(), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("run", "top_p")
    op.drop_column("run", "temperature")
    op.drop_column("run", "usage")
    op.drop_column("run", "truncation_strategy")
    op.drop_column("run", "tool_choice")
    op.drop_column("run", "response_format")
    op.drop_column("run", "max_prompt_tokens")
    op.drop_column("run", "max_completion_tokens")
    op.drop_column("run", "incomplete_details")
    op.drop_column("assistant", "top_p")
    op.drop_column("assistant", "temperature")
    op.drop_column("assistant", "tool_resources")
    op.drop_column("assistant", "response_format")
    # ### end Alembic commands ###