185 lines
9.6 KiB
Python
185 lines
9.6 KiB
Python
"""Migrate agents to orm
|
|
|
|
Revision ID: d05669b60ebe
|
|
Revises: c5d964280dff
|
|
Create Date: 2024-12-12 10:25:31.825635
|
|
|
|
"""
|
|
|
|
from typing import Sequence, Union
|
|
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
from alembic import op
|
|
from letta.settings import settings
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "d05669b60ebe"
|
|
down_revision: Union[str, None] = "c5d964280dff"
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# Skip this migration for SQLite
|
|
if not settings.letta_pg_uri_no_default:
|
|
return
|
|
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
op.create_table(
|
|
"sources_agents",
|
|
sa.Column("agent_id", sa.String(), nullable=False),
|
|
sa.Column("source_id", sa.String(), nullable=False),
|
|
sa.ForeignKeyConstraint(
|
|
["agent_id"],
|
|
["agents.id"],
|
|
),
|
|
sa.ForeignKeyConstraint(
|
|
["source_id"],
|
|
["sources.id"],
|
|
),
|
|
sa.PrimaryKeyConstraint("agent_id", "source_id"),
|
|
)
|
|
op.drop_index("agent_source_mapping_idx_user", table_name="agent_source_mapping")
|
|
op.drop_table("agent_source_mapping")
|
|
op.add_column("agents", sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=True))
|
|
op.add_column("agents", sa.Column("is_deleted", sa.Boolean(), server_default=sa.text("FALSE"), nullable=False))
|
|
op.add_column("agents", sa.Column("_created_by_id", sa.String(), nullable=True))
|
|
op.add_column("agents", sa.Column("_last_updated_by_id", sa.String(), nullable=True))
|
|
op.add_column("agents", sa.Column("organization_id", sa.String(), nullable=True))
|
|
# Populate `organization_id` based on `user_id`
|
|
# Use a raw SQL query to update the organization_id
|
|
op.execute(
|
|
"""
|
|
UPDATE agents
|
|
SET organization_id = users.organization_id
|
|
FROM users
|
|
WHERE agents.user_id = users.id
|
|
"""
|
|
)
|
|
op.alter_column("agents", "organization_id", nullable=False)
|
|
op.alter_column("agents", "name", existing_type=sa.VARCHAR(), nullable=True)
|
|
op.drop_index("agents_idx_user", table_name="agents")
|
|
op.create_unique_constraint("unique_org_agent_name", "agents", ["organization_id", "name"])
|
|
op.create_foreign_key(None, "agents", "organizations", ["organization_id"], ["id"])
|
|
op.drop_column("agents", "tool_names")
|
|
op.drop_column("agents", "user_id")
|
|
op.drop_constraint("agents_tags_organization_id_fkey", "agents_tags", type_="foreignkey")
|
|
op.drop_column("agents_tags", "_created_by_id")
|
|
op.drop_column("agents_tags", "_last_updated_by_id")
|
|
op.drop_column("agents_tags", "updated_at")
|
|
op.drop_column("agents_tags", "id")
|
|
op.drop_column("agents_tags", "is_deleted")
|
|
op.drop_column("agents_tags", "created_at")
|
|
op.drop_column("agents_tags", "organization_id")
|
|
op.create_unique_constraint("unique_agent_block", "blocks_agents", ["agent_id", "block_id"])
|
|
op.drop_constraint("fk_block_id_label", "blocks_agents", type_="foreignkey")
|
|
op.create_foreign_key(
|
|
"fk_block_id_label", "blocks_agents", "block", ["block_id", "block_label"], ["id", "label"], initially="DEFERRED", deferrable=True
|
|
)
|
|
op.drop_column("blocks_agents", "_created_by_id")
|
|
op.drop_column("blocks_agents", "_last_updated_by_id")
|
|
op.drop_column("blocks_agents", "updated_at")
|
|
op.drop_column("blocks_agents", "id")
|
|
op.drop_column("blocks_agents", "is_deleted")
|
|
op.drop_column("blocks_agents", "created_at")
|
|
op.drop_constraint("unique_tool_per_agent", "tools_agents", type_="unique")
|
|
op.create_unique_constraint("unique_agent_tool", "tools_agents", ["agent_id", "tool_id"])
|
|
op.drop_constraint("fk_tool_id", "tools_agents", type_="foreignkey")
|
|
op.drop_constraint("tools_agents_agent_id_fkey", "tools_agents", type_="foreignkey")
|
|
op.create_foreign_key(None, "tools_agents", "tools", ["tool_id"], ["id"], ondelete="CASCADE")
|
|
op.create_foreign_key(None, "tools_agents", "agents", ["agent_id"], ["id"], ondelete="CASCADE")
|
|
op.drop_column("tools_agents", "_created_by_id")
|
|
op.drop_column("tools_agents", "tool_name")
|
|
op.drop_column("tools_agents", "_last_updated_by_id")
|
|
op.drop_column("tools_agents", "updated_at")
|
|
op.drop_column("tools_agents", "id")
|
|
op.drop_column("tools_agents", "is_deleted")
|
|
op.drop_column("tools_agents", "created_at")
|
|
# ### end Alembic commands ###
|
|
|
|
|
|
def downgrade() -> None:
|
|
# Skip this migration for SQLite
|
|
if not settings.letta_pg_uri_no_default:
|
|
return
|
|
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
op.add_column(
|
|
"tools_agents",
|
|
sa.Column("created_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
|
|
)
|
|
op.add_column(
|
|
"tools_agents", sa.Column("is_deleted", sa.BOOLEAN(), server_default=sa.text("false"), autoincrement=False, nullable=False)
|
|
)
|
|
op.add_column("tools_agents", sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False))
|
|
op.add_column(
|
|
"tools_agents",
|
|
sa.Column("updated_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
|
|
)
|
|
op.add_column("tools_agents", sa.Column("_last_updated_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
|
|
op.add_column("tools_agents", sa.Column("tool_name", sa.VARCHAR(), autoincrement=False, nullable=False))
|
|
op.add_column("tools_agents", sa.Column("_created_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
|
|
op.drop_constraint(None, "tools_agents", type_="foreignkey")
|
|
op.drop_constraint(None, "tools_agents", type_="foreignkey")
|
|
op.create_foreign_key("tools_agents_agent_id_fkey", "tools_agents", "agents", ["agent_id"], ["id"])
|
|
op.create_foreign_key("fk_tool_id", "tools_agents", "tools", ["tool_id"], ["id"])
|
|
op.drop_constraint("unique_agent_tool", "tools_agents", type_="unique")
|
|
op.create_unique_constraint("unique_tool_per_agent", "tools_agents", ["agent_id", "tool_name"])
|
|
op.add_column(
|
|
"blocks_agents",
|
|
sa.Column("created_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
|
|
)
|
|
op.add_column(
|
|
"blocks_agents", sa.Column("is_deleted", sa.BOOLEAN(), server_default=sa.text("false"), autoincrement=False, nullable=False)
|
|
)
|
|
op.add_column("blocks_agents", sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False))
|
|
op.add_column(
|
|
"blocks_agents",
|
|
sa.Column("updated_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
|
|
)
|
|
op.add_column("blocks_agents", sa.Column("_last_updated_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
|
|
op.add_column("blocks_agents", sa.Column("_created_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
|
|
op.drop_constraint("fk_block_id_label", "blocks_agents", type_="foreignkey")
|
|
op.create_foreign_key("fk_block_id_label", "blocks_agents", "block", ["block_id", "block_label"], ["id", "label"])
|
|
op.drop_constraint("unique_agent_block", "blocks_agents", type_="unique")
|
|
op.add_column("agents_tags", sa.Column("organization_id", sa.VARCHAR(), autoincrement=False, nullable=False))
|
|
op.add_column(
|
|
"agents_tags",
|
|
sa.Column("created_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
|
|
)
|
|
op.add_column(
|
|
"agents_tags", sa.Column("is_deleted", sa.BOOLEAN(), server_default=sa.text("false"), autoincrement=False, nullable=False)
|
|
)
|
|
op.add_column("agents_tags", sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False))
|
|
op.add_column(
|
|
"agents_tags",
|
|
sa.Column("updated_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
|
|
)
|
|
op.add_column("agents_tags", sa.Column("_last_updated_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
|
|
op.add_column("agents_tags", sa.Column("_created_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
|
|
op.create_foreign_key("agents_tags_organization_id_fkey", "agents_tags", "organizations", ["organization_id"], ["id"])
|
|
op.add_column("agents", sa.Column("user_id", sa.VARCHAR(), autoincrement=False, nullable=False))
|
|
op.add_column("agents", sa.Column("tool_names", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True))
|
|
op.drop_constraint(None, "agents", type_="foreignkey")
|
|
op.drop_constraint("unique_org_agent_name", "agents", type_="unique")
|
|
op.create_index("agents_idx_user", "agents", ["user_id"], unique=False)
|
|
op.alter_column("agents", "name", existing_type=sa.VARCHAR(), nullable=False)
|
|
op.drop_column("agents", "organization_id")
|
|
op.drop_column("agents", "_last_updated_by_id")
|
|
op.drop_column("agents", "_created_by_id")
|
|
op.drop_column("agents", "is_deleted")
|
|
op.drop_column("agents", "updated_at")
|
|
op.create_table(
|
|
"agent_source_mapping",
|
|
sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False),
|
|
sa.Column("user_id", sa.VARCHAR(), autoincrement=False, nullable=False),
|
|
sa.Column("agent_id", sa.VARCHAR(), autoincrement=False, nullable=False),
|
|
sa.Column("source_id", sa.VARCHAR(), autoincrement=False, nullable=False),
|
|
sa.PrimaryKeyConstraint("id", name="agent_source_mapping_pkey"),
|
|
)
|
|
op.create_index("agent_source_mapping_idx_user", "agent_source_mapping", ["user_id", "agent_id", "source_id"], unique=False)
|
|
op.drop_table("sources_agents")
|
|
# ### end Alembic commands ###
|