"""Add task_id, conversation_id to models

Revision ID: f39011d0e6df
Revises: bf46ba93762f
Create Date: 2023-06-30 07:36:18.677148

"""

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "f39011d0e6df"
down_revision = "bf46ba93762f"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_unique_constraint(
        None,
        "inference_prompt_contents",
        ["inference_prompt_id"],
    )
    op.create_unique_constraint(
        None,
        "inference_response_contents",
        ["inference_response_id"],
    )
    op.add_column("inferences", sa.Column("task_id", sa.String(), nullable=True))
    op.add_column(
        "inferences",
        sa.Column("conversation_id", sa.String(), nullable=True),
    )
    op.create_index(
        op.f("ix_inferences_conversation_id"),
        "inferences",
        ["conversation_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_inferences_task_id"),
        "inferences",
        ["task_id"],
        unique=False,
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_inferences_task_id"), table_name="inferences")
    op.drop_index(op.f("ix_inferences_conversation_id"), table_name="inferences")
    op.drop_column("inferences", "conversation_id")
    op.drop_column("inferences", "task_id")
    op.drop_constraint(None, "inference_response_contents", type_="unique")
    op.drop_constraint(None, "inference_prompt_contents", type_="unique")
    # ### end Alembic commands ###
