AgentEvalTool/migrations/versions/a64b2f8c9d10_add_model_metadata.py

66 lines
2.6 KiB
Python

"""add model metadata
Revision ID: a64b2f8c9d10
Revises: 8e91c70a5d3b
Create Date: 2026-07-17
"""
from typing import Sequence, Union
import sqlalchemy as sa
import sqlmodel
from alembic import op
revision: str = "a64b2f8c9d10"
down_revision: Union[str, Sequence[str], None] = "8e91c70a5d3b"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
with op.batch_alter_table("model_configs", schema=None) as batch_op:
batch_op.add_column(
sa.Column("vendor_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False, server_default="")
)
batch_op.add_column(
sa.Column(
"input_modalities",
sqlmodel.sql.sqltypes.AutoString(),
nullable=False,
server_default='["text"]',
)
)
batch_op.add_column(
sa.Column(
"output_modalities",
sqlmodel.sql.sqltypes.AutoString(),
nullable=False,
server_default='["text"]',
)
)
batch_op.add_column(sa.Column("context_window", sa.Integer(), nullable=True))
batch_op.add_column(sa.Column("max_output_tokens", sa.Integer(), nullable=True))
batch_op.add_column(sa.Column("supports_streaming", sa.Boolean(), nullable=False, server_default=sa.false()))
batch_op.add_column(sa.Column("supports_tool_calling", sa.Boolean(), nullable=False, server_default=sa.false()))
batch_op.add_column(
sa.Column("supports_structured_output", sa.Boolean(), nullable=False, server_default=sa.false())
)
batch_op.add_column(sa.Column("supports_reasoning", sa.Boolean(), nullable=False, server_default=sa.false()))
batch_op.add_column(sa.Column("region", sqlmodel.sql.sqltypes.AutoString(), nullable=False, server_default=""))
batch_op.add_column(sa.Column("documentation_url", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
def downgrade() -> None:
with op.batch_alter_table("model_configs", schema=None) as batch_op:
batch_op.drop_column("documentation_url")
batch_op.drop_column("region")
batch_op.drop_column("supports_reasoning")
batch_op.drop_column("supports_structured_output")
batch_op.drop_column("supports_tool_calling")
batch_op.drop_column("supports_streaming")
batch_op.drop_column("max_output_tokens")
batch_op.drop_column("context_window")
batch_op.drop_column("output_modalities")
batch_op.drop_column("input_modalities")
batch_op.drop_column("vendor_name")