diff --git a/alembic/versions/25fc99e97839_fix_alembic_check_warnings.py b/alembic/versions/25fc99e97839_fix_alembic_check_warnings.py new file mode 100644 index 00000000..47566a8a --- /dev/null +++ b/alembic/versions/25fc99e97839_fix_alembic_check_warnings.py @@ -0,0 +1,43 @@ +"""Remove job_usage_statistics indices and update job_messages + +Revision ID: 25fc99e97839 +Revises: f595e0e8013e +Create Date: 2025-01-16 16:48:21.000000 + +""" + +from typing import Sequence, Union + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "25fc99e97839" +down_revision: Union[str, None] = "f595e0e8013e" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Remove indices from job_messages + op.drop_index("ix_job_messages_created_at", table_name="job_messages") + op.drop_index("ix_job_messages_job_id", table_name="job_messages") + + # Remove indices from job_usage_statistics + op.drop_index("ix_job_usage_statistics_created_at", table_name="job_usage_statistics") + op.drop_index("ix_job_usage_statistics_job_id", table_name="job_usage_statistics") + + # Add foreign key constraint for message_id + op.create_foreign_key("fk_job_messages_message_id", "job_messages", "messages", ["message_id"], ["id"], ondelete="CASCADE") + + +def downgrade() -> None: + # Remove the foreign key constraint + op.drop_constraint("fk_job_messages_message_id", "job_messages", type_="foreignkey") + + # Recreate indices for job_messages + op.create_index("ix_job_messages_job_id", "job_messages", ["job_id"]) + op.create_index("ix_job_messages_created_at", "job_messages", ["created_at"]) + + # Recreate indices for job_usage_statistics + op.create_index("ix_job_usage_statistics_job_id", "job_usage_statistics", ["job_id"]) + op.create_index("ix_job_usage_statistics_created_at", "job_usage_statistics", ["created_at"]) diff --git a/letta/orm/__init__.py b/letta/orm/__init__.py index 79185aa4..74cfe0c4 100644 --- a/letta/orm/__init__.py +++ b/letta/orm/__init__.py @@ -6,6 +6,7 @@ from letta.orm.blocks_agents import BlocksAgents from letta.orm.file import FileMetadata from letta.orm.job import Job from letta.orm.job_messages import JobMessage +from letta.orm.job_usage_statistics import JobUsageStatistics from letta.orm.message import Message from letta.orm.organization import Organization from letta.orm.passage import AgentPassage, BasePassage, SourcePassage