67 lines
4.2 KiB
Python
67 lines
4.2 KiB
Python
"""Persist chat state, indexing tasks and fixed batch membership.
|
|
|
|
Revision ID: 0003
|
|
Revises: 0002
|
|
"""
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
revision = "0003"
|
|
down_revision = "0002"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def upgrade():
|
|
inspector = sa.inspect(op.get_bind())
|
|
tables = set(inspector.get_table_names())
|
|
if "project_ids" not in {column["name"] for column in inspector.get_columns("scoring_batches")}:
|
|
op.add_column("scoring_batches", sa.Column("project_ids", sa.JSON(), nullable=True))
|
|
if "index_tasks" not in tables:
|
|
op.create_table("index_tasks",
|
|
sa.Column("id", sa.Integer(), primary_key=True),
|
|
sa.Column("course_version_id", sa.Integer(), sa.ForeignKey("course_versions.id"), nullable=False),
|
|
sa.Column("submission_file_id", sa.Integer(), sa.ForeignKey("submission_files.id")),
|
|
sa.Column("kind", sa.String(30), nullable=False), sa.Column("status", sa.String(30), nullable=False),
|
|
sa.Column("result", sa.JSON(), nullable=False), sa.Column("error", sa.Text()),
|
|
sa.Column("created_at", sa.DateTime(timezone=True)), sa.Column("updated_at", sa.DateTime(timezone=True)))
|
|
if "conversations" not in tables:
|
|
op.create_table("conversations",
|
|
sa.Column("id", sa.Integer(), primary_key=True), sa.Column("title", sa.String(200), nullable=False),
|
|
sa.Column("course_version_id", sa.Integer(), sa.ForeignKey("course_versions.id")),
|
|
sa.Column("batch_id", sa.Integer(), sa.ForeignKey("scoring_batches.id")),
|
|
sa.Column("project_id", sa.Integer(), sa.ForeignKey("projects.id")),
|
|
sa.Column("created_at", sa.DateTime(timezone=True)), sa.Column("updated_at", sa.DateTime(timezone=True)))
|
|
if "chat_messages" not in tables:
|
|
op.create_table("chat_messages",
|
|
sa.Column("id", sa.Integer(), primary_key=True),
|
|
sa.Column("conversation_id", sa.Integer(), sa.ForeignKey("conversations.id"), nullable=False),
|
|
sa.Column("role", sa.String(20), nullable=False), sa.Column("content", sa.Text(), nullable=False),
|
|
sa.Column("citations", sa.JSON(), nullable=False), sa.Column("tool_calls", sa.JSON(), nullable=False),
|
|
sa.Column("created_at", sa.DateTime(timezone=True)))
|
|
if "agent_runs" not in tables:
|
|
op.create_table("agent_runs", sa.Column("id", sa.Integer(), primary_key=True),
|
|
sa.Column("conversation_id", sa.Integer(), sa.ForeignKey("conversations.id"), nullable=False),
|
|
sa.Column("status", sa.String(30), nullable=False), sa.Column("error", sa.Text()),
|
|
sa.Column("created_at", sa.DateTime(timezone=True)), sa.Column("finished_at", sa.DateTime(timezone=True)))
|
|
if "pending_actions" not in tables:
|
|
op.create_table("pending_actions", sa.Column("id", sa.Integer(), primary_key=True),
|
|
sa.Column("conversation_id", sa.Integer(), sa.ForeignKey("conversations.id"), nullable=False),
|
|
sa.Column("kind", sa.String(40), nullable=False), sa.Column("payload", sa.JSON(), nullable=False),
|
|
sa.Column("preview", sa.JSON(), nullable=False), sa.Column("expected", sa.JSON(), nullable=False),
|
|
sa.Column("idempotency_key", sa.String(64), nullable=False, unique=True),
|
|
sa.Column("status", sa.String(30), nullable=False), sa.Column("result", sa.JSON()),
|
|
sa.Column("created_at", sa.DateTime(timezone=True)), sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False))
|
|
if "ix_chat_messages_conversation_id" not in {index["name"] for index in sa.inspect(op.get_bind()).get_indexes("chat_messages")}:
|
|
op.create_index("ix_chat_messages_conversation_id", "chat_messages", ["conversation_id"])
|
|
if "ix_pending_actions_conversation_id" not in {index["name"] for index in sa.inspect(op.get_bind()).get_indexes("pending_actions")}:
|
|
op.create_index("ix_pending_actions_conversation_id", "pending_actions", ["conversation_id"])
|
|
|
|
|
|
def downgrade():
|
|
op.drop_index("ix_pending_actions_conversation_id", table_name="pending_actions")
|
|
op.drop_index("ix_chat_messages_conversation_id", table_name="chat_messages")
|
|
for table in ("pending_actions", "agent_runs", "chat_messages", "conversations", "index_tasks"):
|
|
op.drop_table(table)
|
|
op.drop_column("scoring_batches", "project_ids")
|