"""baseline Revision ID: db67202be996 Revises: Create Date: 2026-03-24 23:17:38.516658 """ from typing import Sequence, Union import sqlalchemy as sa from alembic import op from migrations.utils.safe_ddl import ( create_index_if_not_exists, drop_index_if_exists, drop_table_if_exists, has_table, ) # revision identifiers, used by Alembic. revision: str = "db67202be996" down_revision: Union[str, None] = None branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None _STATUS_ENUM = sa.Enum( "deleted", "active", "queued", "in_progress", "pending_action", "completed", "failed", "cancelling", "cancelled", "pending", "processing", "expired", "retrying", "offline", name="statusenum", ) def upgrade() -> None: """Upgrade schema.""" # ------------------------------------------------------------------------- # base_models # ------------------------------------------------------------------------- if not has_table("base_models"): op.create_table( "base_models", sa.Column("id", sa.String(length=128), nullable=False), sa.Column("name", sa.String(length=128), nullable=False), sa.Column("family", sa.String(length=64), nullable=True), sa.Column("parameter_count", sa.String(length=32), nullable=True), sa.Column("is_multimodal", sa.Boolean(), nullable=True), sa.Column("created_at", sa.BigInteger(), nullable=True), sa.Column( "endpoint", sa.String(length=256), nullable=True, comment="HuggingFace model path or custom inference endpoint identifier.", ), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_base_models_id"), "base_models", ["id"], unique=False ) # ------------------------------------------------------------------------- # compute_nodes # ------------------------------------------------------------------------- if not has_table("compute_nodes"): op.create_table( "compute_nodes", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("hostname", sa.String(length=128), nullable=False), sa.Column("ip_address", sa.String(length=45), nullable=True), sa.Column("gpu_model", sa.String(length=128), nullable=True), sa.Column("total_vram_gb", sa.Float(), nullable=True), sa.Column("current_vram_usage_gb", sa.Float(), nullable=True), sa.Column("status", _STATUS_ENUM, nullable=True), sa.Column("last_heartbeat", sa.BigInteger(), nullable=True), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_compute_nodes_id"), "compute_nodes", ["id"], unique=False ) # ------------------------------------------------------------------------- # messages # ------------------------------------------------------------------------- if not has_table("messages"): op.create_table( "messages", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("assistant_id", sa.String(length=64), nullable=True), sa.Column("attachments", sa.JSON(), nullable=True), sa.Column("completed_at", sa.Integer(), nullable=True), sa.Column("content", sa.Text(length=4294967295), nullable=False), sa.Column( "reasoning", sa.Text(length=4294967295), nullable=True, comment="Stores the internal 'thinking' or reasoning tokens from the model.", ), sa.Column( "tool_calls", sa.JSON(), nullable=True, comment="Stores the native JSON array of tool calls generated by the assistant.", ), sa.Column( "tool_call_id", sa.String(length=64), nullable=True, comment="For messages with role='tool', this links back to the specific tool_call_id.", ), sa.Column("created_at", sa.Integer(), nullable=False), sa.Column("incomplete_at", sa.Integer(), nullable=True), sa.Column("incomplete_details", sa.JSON(), nullable=True), sa.Column("meta_data", sa.JSON(), nullable=False), sa.Column("object", sa.String(length=64), nullable=False), sa.Column("role", sa.String(length=32), nullable=False), sa.Column("run_id", sa.String(length=64), nullable=True), sa.Column("tool_id", sa.String(length=64), nullable=True), sa.Column("status", sa.String(length=32), nullable=True), sa.Column("thread_id", sa.String(length=64), nullable=False), sa.Column("sender_id", sa.String(length=64), nullable=True), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_messages_assistant_id"), "messages", ["assistant_id"], unique=False ) create_index_if_not_exists(op.f("ix_messages_id"), "messages", ["id"], unique=False) # ------------------------------------------------------------------------- # users # ------------------------------------------------------------------------- if not has_table("users"): op.create_table( "users", sa.Column( "id", sa.String(length=64), nullable=False, comment="Internal unique identifier for the user (e.g., user_...)", ), sa.Column("created_at", sa.DateTime(), nullable=False), sa.Column("updated_at", sa.DateTime(), nullable=False), sa.Column( "is_admin", sa.Boolean(), server_default="0", nullable=False, comment="Flag indicating administrative privileges", ), sa.Column( "email", sa.String(length=255), nullable=True, comment="Primary email address, potentially verified by OAuth provider", ), sa.Column( "email_verified", sa.Boolean(), nullable=True, comment="Whether the email address has been verified", ), sa.Column( "full_name", sa.String(length=255), nullable=True, comment="User's full display name", ), sa.Column( "given_name", sa.String(length=128), nullable=True, comment="First name" ), sa.Column( "family_name", sa.String(length=128), nullable=True, comment="Last name" ), sa.Column( "picture_url", sa.Text(), nullable=True, comment="URL to the user's profile picture", ), sa.Column( "oauth_provider", sa.String(length=50), nullable=True, comment="Name of the OAuth provider (e.g., 'google', 'github', 'local')", ), sa.Column( "provider_user_id", sa.String(length=255), nullable=True, comment="The unique ID assigned by the OAuth provider", ), sa.PrimaryKeyConstraint("id"), sa.UniqueConstraint( "oauth_provider", "provider_user_id", name="uq_user_oauth_provider_id" ), ) create_index_if_not_exists("idx_user_email", "users", ["email"], unique=False) create_index_if_not_exists("idx_user_is_admin", "users", ["is_admin"], unique=False) create_index_if_not_exists(op.f("ix_users_email"), "users", ["email"], unique=True) create_index_if_not_exists(op.f("ix_users_id"), "users", ["id"], unique=False) create_index_if_not_exists( op.f("ix_users_oauth_provider"), "users", ["oauth_provider"], unique=False ) create_index_if_not_exists( op.f("ix_users_provider_user_id"), "users", ["provider_user_id"], unique=False ) # ------------------------------------------------------------------------- # api_keys # ------------------------------------------------------------------------- if not has_table("api_keys"): op.create_table( "api_keys", sa.Column("id", sa.Integer(), nullable=False), sa.Column("key_name", sa.String(length=100), nullable=True), sa.Column("hashed_key", sa.String(length=255), nullable=False), sa.Column("prefix", sa.String(length=8), nullable=False), sa.Column("user_id", sa.String(length=64), nullable=False), sa.Column("created_at", sa.DateTime(), nullable=True), sa.Column("expires_at", sa.DateTime(), nullable=True), sa.Column("last_used_at", sa.DateTime(), nullable=True), sa.Column("is_active", sa.Boolean(), nullable=False), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), sa.UniqueConstraint("prefix"), ) create_index_if_not_exists( "idx_apikey_user_id_active", "api_keys", ["user_id", "is_active"], unique=False ) create_index_if_not_exists( op.f("ix_api_keys_hashed_key"), "api_keys", ["hashed_key"], unique=True ) create_index_if_not_exists(op.f("ix_api_keys_id"), "api_keys", ["id"], unique=False) create_index_if_not_exists( op.f("ix_api_keys_user_id"), "api_keys", ["user_id"], unique=False ) # ------------------------------------------------------------------------- # assistants # ------------------------------------------------------------------------- if not has_table("assistants"): op.create_table( "assistants", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("object", sa.String(length=64), nullable=False), sa.Column("created_at", sa.Integer(), nullable=False), sa.Column("name", sa.String(length=128), nullable=False), sa.Column("description", sa.String(length=256), nullable=True), sa.Column("model", sa.String(length=64), nullable=True), sa.Column("instructions", sa.Text(), nullable=True), sa.Column("tool_configs", sa.JSON(), nullable=True), sa.Column("meta_data", sa.JSON(), nullable=True), sa.Column("top_p", sa.Integer(), nullable=True), sa.Column("temperature", sa.Integer(), nullable=True), sa.Column("response_format", sa.String(length=64), nullable=True), sa.Column( "tool_resources", sa.JSON(), nullable=True, comment='Resource map keyed by tool type, e.g. {"file_search": {"vector_store_ids": ["vs_123","vs_456"]}}', ), sa.Column( "owner_id", sa.String(length=64), nullable=True, comment="Canonical owner of this assistant. Primary key for row-level access filtering. Separate from the many-to-many (user_assistants) which handles sharing.", ), sa.Column( "max_turns", sa.Integer(), server_default="1", nullable=True, comment="Max number of iterative loops for Level 3 agency. 1 = Standard Level 2 (ReAct).", ), sa.Column( "agent_mode", sa.Boolean(), server_default="0", nullable=False, comment="False = Standard (Level 2), True = Autonomous (Level 3).", ), sa.Column( "web_access", sa.Boolean(), server_default="0", nullable=False, comment="Enable live web search and browsing capabilities.", ), sa.Column( "deep_research", sa.Boolean(), server_default="0", nullable=False, comment="Enable deep research capabilities.", ), sa.Column( "engineer", sa.Boolean(), server_default="0", nullable=False, comment="Enable network engineering capabilities and inventory map access.", ), sa.Column( "decision_telemetry", sa.Boolean(), server_default="0", nullable=False, comment="If True, captures detailed reasoning payloads and confidence scores.", ), sa.Column( "deleted_at", sa.Integer(), nullable=True, comment="Unix timestamp of soft-deletion. If present, entity is in 'Recycle Bin'.", ), sa.ForeignKeyConstraint(["owner_id"], ["users.id"], ondelete="SET NULL"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_assistants_id"), "assistants", ["id"], unique=False ) create_index_if_not_exists( op.f("ix_assistants_owner_id"), "assistants", ["owner_id"], unique=False ) # ------------------------------------------------------------------------- # audit_logs # ------------------------------------------------------------------------- if not has_table("audit_logs"): op.create_table( "audit_logs", sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False), sa.Column("user_id", sa.String(length=64), nullable=True), sa.Column( "action", sa.String(length=32), nullable=False, comment="e.g. CREATE, UPDATE, DELETE, HARD_DELETE, ERASE", ), sa.Column( "entity_type", sa.String(length=64), nullable=False, comment="e.g. Assistant, User, Thread", ), sa.Column("entity_id", sa.String(length=64), nullable=False), sa.Column("timestamp", sa.DateTime(), nullable=False), sa.Column("ip_address", sa.String(length=45), nullable=True), sa.Column( "details", sa.JSON(), nullable=True, comment="Stores before/after state or reasoning for action", ), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="SET NULL"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_audit_logs_action"), "audit_logs", ["action"], unique=False ) create_index_if_not_exists( op.f("ix_audit_logs_entity_id"), "audit_logs", ["entity_id"], unique=False ) create_index_if_not_exists( op.f("ix_audit_logs_entity_type"), "audit_logs", ["entity_type"], unique=False ) create_index_if_not_exists( op.f("ix_audit_logs_id"), "audit_logs", ["id"], unique=False ) create_index_if_not_exists( op.f("ix_audit_logs_timestamp"), "audit_logs", ["timestamp"], unique=False ) create_index_if_not_exists( op.f("ix_audit_logs_user_id"), "audit_logs", ["user_id"], unique=False ) # ------------------------------------------------------------------------- # batfish_snapshots # ------------------------------------------------------------------------- if not has_table("batfish_snapshots"): op.create_table( "batfish_snapshots", sa.Column( "id", sa.String(length=64), nullable=False, comment="Opaque snapshot ID returned to caller e.g. snap_abc123", ), sa.Column( "snapshot_name", sa.String(length=128), nullable=False, comment="Caller-supplied label e.g. 'incident_001'", ), sa.Column( "snapshot_key", sa.String(length=256), nullable=False, comment="Namespaced isolation key: {user_id}_{id}", ), sa.Column("user_id", sa.String(length=64), nullable=False), sa.Column("configs_root", sa.String(length=512), nullable=True), sa.Column("device_count", sa.Integer(), nullable=False), sa.Column( "devices", sa.JSON(), nullable=False, comment="List of hostnames ingested into this snapshot", ), sa.Column("status", _STATUS_ENUM, nullable=False), sa.Column("error_message", sa.Text(), nullable=True), sa.Column("created_at", sa.BigInteger(), nullable=False), sa.Column("updated_at", sa.BigInteger(), nullable=False), sa.Column("last_ingested_at", sa.BigInteger(), nullable=True), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), sa.UniqueConstraint( "user_id", "snapshot_name", name="uq_batfish_user_snapshot_name" ), ) create_index_if_not_exists( "idx_batfish_status", "batfish_snapshots", ["status"], unique=False ) create_index_if_not_exists( "idx_batfish_user_id", "batfish_snapshots", ["user_id"], unique=False ) create_index_if_not_exists( op.f("ix_batfish_snapshots_id"), "batfish_snapshots", ["id"], unique=False ) create_index_if_not_exists( op.f("ix_batfish_snapshots_snapshot_key"), "batfish_snapshots", ["snapshot_key"], unique=True, ) create_index_if_not_exists( op.f("ix_batfish_snapshots_user_id"), "batfish_snapshots", ["user_id"], unique=False, ) # ------------------------------------------------------------------------- # datasets # ------------------------------------------------------------------------- if not has_table("datasets"): op.create_table( "datasets", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("user_id", sa.String(length=64), nullable=False), sa.Column("name", sa.String(length=128), nullable=False), sa.Column("description", sa.Text(), nullable=True), sa.Column("format", sa.String(length=32), nullable=False), sa.Column("file_id", sa.String(length=64), nullable=False), sa.Column("storage_path", sa.String(length=512), nullable=True), sa.Column("train_samples", sa.Integer(), nullable=True), sa.Column("eval_samples", sa.Integer(), nullable=True), sa.Column("config", sa.JSON(), nullable=True), sa.Column("status", _STATUS_ENUM, nullable=False), sa.Column("created_at", sa.BigInteger(), nullable=False), sa.Column("updated_at", sa.BigInteger(), nullable=False), sa.Column("deleted_at", sa.Integer(), nullable=True), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_datasets_deleted_at"), "datasets", ["deleted_at"], unique=False ) create_index_if_not_exists( op.f("ix_datasets_file_id"), "datasets", ["file_id"], unique=False ) create_index_if_not_exists(op.f("ix_datasets_id"), "datasets", ["id"], unique=False) create_index_if_not_exists( op.f("ix_datasets_user_id"), "datasets", ["user_id"], unique=False ) # ------------------------------------------------------------------------- # files # ------------------------------------------------------------------------- if not has_table("files"): op.create_table( "files", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("object", sa.String(length=64), nullable=False), sa.Column("bytes", sa.Integer(), nullable=False), sa.Column("created_at", sa.DateTime(), nullable=False), sa.Column("expires_at", sa.DateTime(), nullable=True), sa.Column("filename", sa.String(length=256), nullable=False), sa.Column("purpose", sa.String(length=64), nullable=False), sa.Column("mime_type", sa.String(length=255), nullable=True), sa.Column( "deleted_at", sa.Integer(), nullable=True, comment="Unix timestamp of soft-deletion. Non-null = file is in Recycle Bin.", ), sa.Column("user_id", sa.String(length=64), nullable=False), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_files_deleted_at"), "files", ["deleted_at"], unique=False ) create_index_if_not_exists(op.f("ix_files_id"), "files", ["id"], unique=False) # ------------------------------------------------------------------------- # runs # ------------------------------------------------------------------------- if not has_table("runs"): op.create_table( "runs", sa.Column("user_id", sa.String(length=64), nullable=True), sa.Column("id", sa.String(length=64), nullable=False), sa.Column("assistant_id", sa.String(length=64), nullable=False), sa.Column("cancelled_at", sa.Integer(), nullable=True), sa.Column("completed_at", sa.Integer(), nullable=True), sa.Column("created_at", sa.Integer(), nullable=True), sa.Column("expires_at", sa.Integer(), nullable=True), sa.Column("failed_at", sa.Integer(), nullable=True), sa.Column("started_at", sa.Integer(), nullable=True), sa.Column("incomplete_details", sa.String(length=256), nullable=True), sa.Column("instructions", sa.Text(), nullable=True), sa.Column("last_error", sa.String(length=256), nullable=True), sa.Column("max_completion_tokens", sa.Integer(), nullable=True), sa.Column("max_prompt_tokens", sa.Integer(), nullable=True), sa.Column("meta_data", sa.JSON(), nullable=True), sa.Column("model", sa.String(length=64), nullable=True), sa.Column("object", sa.String(length=64), nullable=False), sa.Column("parallel_tool_calls", sa.Boolean(), nullable=True), sa.Column("required_action", sa.String(length=256), nullable=True), sa.Column("response_format", sa.String(length=64), nullable=True), sa.Column("status", _STATUS_ENUM, nullable=False), sa.Column("thread_id", sa.String(length=64), nullable=False), sa.Column("tool_choice", sa.String(length=64), nullable=True), sa.Column("tools", sa.JSON(), nullable=True), sa.Column( "current_turn", sa.Integer(), server_default="0", nullable=True, comment="The current iteration count of the autonomous loop.", ), sa.Column( "max_turns", sa.Integer(), server_default="1", nullable=True, comment="Snapshot of Assistant.max_turns at run creation.", ), sa.Column( "agent_mode", sa.String(length=32), server_default="standard", nullable=True, comment="Determines execution logic: 'standard' (Level 2) or 'autonomous' (Level 3).", ), sa.Column( "truncation_strategy", sa.String(length=16), server_default="auto", nullable=True, ), sa.Column("usage", sa.JSON(), nullable=True), sa.Column("temperature", sa.Integer(), nullable=True), sa.Column("top_p", sa.Integer(), nullable=True), sa.Column("tool_resources", sa.JSON(), nullable=True), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_runs_user_id"), "runs", ["user_id"], unique=False ) # ------------------------------------------------------------------------- # sandboxes # ------------------------------------------------------------------------- if not has_table("sandboxes"): op.create_table( "sandboxes", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("user_id", sa.String(length=64), nullable=False), sa.Column("name", sa.String(length=128), nullable=False), sa.Column("created_at", sa.DateTime(), nullable=True), sa.Column("status", sa.String(length=32), nullable=False), sa.Column("config", sa.JSON(), nullable=True), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_sandboxes_id"), "sandboxes", ["id"], unique=False ) # ------------------------------------------------------------------------- # threads # ------------------------------------------------------------------------- if not has_table("threads"): op.create_table( "threads", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("created_at", sa.Integer(), nullable=False), sa.Column("meta_data", sa.JSON(), nullable=False), sa.Column("object", sa.String(length=64), nullable=False), sa.Column("tool_resources", sa.JSON(), nullable=False), sa.Column( "owner_id", sa.String(length=64), nullable=True, comment="Canonical creator/owner of this thread. Used for row-level access control.", ), sa.ForeignKeyConstraint(["owner_id"], ["users.id"], ondelete="SET NULL"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists(op.f("ix_threads_id"), "threads", ["id"], unique=False) create_index_if_not_exists( op.f("ix_threads_owner_id"), "threads", ["owner_id"], unique=False ) # ------------------------------------------------------------------------- # vector_stores # ------------------------------------------------------------------------- if not has_table("vector_stores"): op.create_table( "vector_stores", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("name", sa.String(length=128), nullable=False), sa.Column("user_id", sa.String(length=64), nullable=False), sa.Column("collection_name", sa.String(length=128), nullable=False), sa.Column("vector_size", sa.Integer(), nullable=False), sa.Column("distance_metric", sa.String(length=32), nullable=False), sa.Column("created_at", sa.BigInteger(), nullable=True), sa.Column("updated_at", sa.BigInteger(), nullable=True), sa.Column("status", _STATUS_ENUM, nullable=False), sa.Column("config", sa.JSON(), nullable=True), sa.Column("file_count", sa.Integer(), nullable=False), sa.Column( "deleted_at", sa.Integer(), nullable=True, comment="Unix timestamp of soft-deletion. Non-null = store is in Recycle Bin.", ), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), sa.UniqueConstraint("collection_name"), ) create_index_if_not_exists( op.f("ix_vector_stores_deleted_at"), "vector_stores", ["deleted_at"], unique=False, ) create_index_if_not_exists( op.f("ix_vector_stores_id"), "vector_stores", ["id"], unique=False ) # ------------------------------------------------------------------------- # actions # ------------------------------------------------------------------------- if not has_table("actions"): op.create_table( "actions", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("run_id", sa.String(length=64), nullable=True), sa.Column( "tool_call_id", sa.String(length=64), nullable=True, comment="The unique ID linking this action to a specific LLM tool request.", ), sa.Column( "tool_name", sa.String(length=64), nullable=True, comment="The name of the function/tool to be executed.", ), sa.Column( "turn_index", sa.Integer(), nullable=True, comment="The iteration of the autonomous loop.", ), sa.Column( "decision_payload", sa.JSON(), nullable=True, comment="The full structured reasoning object (reason, policy, etc) preceding the tool call.", ), sa.Column( "confidence_score", sa.Float(), nullable=True, comment="Extracted confidence score (0.0-1.0) to allow fast querying of 'uncertain' agent actions.", ), sa.Column("triggered_at", sa.DateTime(), nullable=True), sa.Column("expires_at", sa.DateTime(), nullable=True), sa.Column("is_processed", sa.Boolean(), nullable=True), sa.Column("processed_at", sa.DateTime(), nullable=True), sa.Column("status", sa.String(length=64), nullable=True), sa.Column("function_args", sa.JSON(), nullable=True), sa.Column("result", sa.JSON(), nullable=True), sa.ForeignKeyConstraint(["run_id"], ["runs.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_actions_confidence_score"), "actions", ["confidence_score"], unique=False, ) create_index_if_not_exists(op.f("ix_actions_id"), "actions", ["id"], unique=False) create_index_if_not_exists( op.f("ix_actions_tool_call_id"), "actions", ["tool_call_id"], unique=False ) # ------------------------------------------------------------------------- # file_storage # ------------------------------------------------------------------------- if not has_table("file_storage"): op.create_table( "file_storage", sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), sa.Column("file_id", sa.String(length=64), nullable=False), sa.Column( "storage_system", sa.String(length=64), nullable=False, comment="Storage system type (samba, s3, etc.)", ), sa.Column( "storage_path", sa.String(length=512), nullable=True, comment="Path to file in storage system (relative to share root)", ), sa.Column( "is_primary", sa.Boolean(), nullable=True, comment="Indicates if this is the primary storage location", ), sa.Column( "created_at", sa.DateTime(), nullable=False, comment="When this storage entry was created", ), sa.ForeignKeyConstraint(["file_id"], ["files.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( "idx_file_storage_file_id", "file_storage", ["file_id"], unique=False ) # ------------------------------------------------------------------------- # thread_participants # ------------------------------------------------------------------------- if not has_table("thread_participants"): op.create_table( "thread_participants", sa.Column("thread_id", sa.String(length=64), nullable=False), sa.Column("user_id", sa.String(length=64), nullable=False), sa.ForeignKeyConstraint(["thread_id"], ["threads.id"], ondelete="CASCADE"), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("thread_id", "user_id"), ) # ------------------------------------------------------------------------- # training_jobs # ------------------------------------------------------------------------- if not has_table("training_jobs"): op.create_table( "training_jobs", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("user_id", sa.String(length=64), nullable=False), sa.Column("dataset_id", sa.String(length=64), nullable=True), sa.Column("base_model", sa.String(length=256), nullable=False), sa.Column("framework", sa.String(length=32), nullable=False), sa.Column("config", sa.JSON(), nullable=True), sa.Column("status", _STATUS_ENUM, nullable=False), sa.Column("created_at", sa.BigInteger(), nullable=False), sa.Column("updated_at", sa.BigInteger(), nullable=False), sa.Column("deleted_at", sa.Integer(), nullable=True), sa.Column("started_at", sa.BigInteger(), nullable=True), sa.Column("completed_at", sa.BigInteger(), nullable=True), sa.Column("failed_at", sa.BigInteger(), nullable=True), sa.Column("last_error", sa.Text(), nullable=True), sa.Column("metrics", sa.JSON(), nullable=True), sa.Column("output_path", sa.String(length=512), nullable=True), sa.Column("node_id", sa.String(length=64), nullable=True), sa.ForeignKeyConstraint( ["dataset_id"], ["datasets.id"], ondelete="SET NULL" ), sa.ForeignKeyConstraint( ["node_id"], ["compute_nodes.id"], ondelete="SET NULL" ), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( "idx_trainingjob_status", "training_jobs", ["status"], unique=False ) create_index_if_not_exists( "idx_trainingjob_user_id", "training_jobs", ["user_id"], unique=False ) create_index_if_not_exists( op.f("ix_training_jobs_dataset_id"), "training_jobs", ["dataset_id"], unique=False, ) create_index_if_not_exists( op.f("ix_training_jobs_deleted_at"), "training_jobs", ["deleted_at"], unique=False, ) create_index_if_not_exists( op.f("ix_training_jobs_id"), "training_jobs", ["id"], unique=False ) create_index_if_not_exists( op.f("ix_training_jobs_user_id"), "training_jobs", ["user_id"], unique=False ) # ------------------------------------------------------------------------- # user_assistants # ------------------------------------------------------------------------- if not has_table("user_assistants"): op.create_table( "user_assistants", sa.Column("user_id", sa.String(length=64), nullable=False), sa.Column("assistant_id", sa.String(length=64), nullable=False), sa.ForeignKeyConstraint( ["assistant_id"], ["assistants.id"], ondelete="CASCADE" ), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("user_id", "assistant_id"), ) # ------------------------------------------------------------------------- # vector_store_files # ------------------------------------------------------------------------- if not has_table("vector_store_files"): op.create_table( "vector_store_files", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("vector_store_id", sa.String(length=64), nullable=False), sa.Column("file_name", sa.String(length=256), nullable=False), sa.Column("file_path", sa.String(length=1024), nullable=False), sa.Column("processed_at", sa.Integer(), nullable=True), sa.Column("status", _STATUS_ENUM, nullable=True), sa.Column("error_message", sa.Text(), nullable=True), sa.Column("meta_data", sa.JSON(), nullable=True), sa.ForeignKeyConstraint( ["vector_store_id"], ["vector_stores.id"], ondelete="CASCADE" ), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_vector_store_files_id"), "vector_store_files", ["id"], unique=False ) # ------------------------------------------------------------------------- # fine_tuned_models # ------------------------------------------------------------------------- if not has_table("fine_tuned_models"): op.create_table( "fine_tuned_models", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("user_id", sa.String(length=64), nullable=False), sa.Column("training_job_id", sa.String(length=64), nullable=True), sa.Column("name", sa.String(length=128), nullable=False), sa.Column("description", sa.Text(), nullable=True), sa.Column("base_model", sa.String(length=256), nullable=False), sa.Column("hf_repo", sa.String(length=256), nullable=True), sa.Column("storage_path", sa.String(length=512), nullable=True), sa.Column("is_active", sa.Boolean(), nullable=False), sa.Column("vllm_model_id", sa.String(length=256), nullable=True), sa.Column("status", _STATUS_ENUM, nullable=False), sa.Column("created_at", sa.BigInteger(), nullable=False), sa.Column("updated_at", sa.BigInteger(), nullable=False), sa.Column("deleted_at", sa.Integer(), nullable=True), sa.Column("node_id", sa.String(length=64), nullable=True), sa.ForeignKeyConstraint( ["node_id"], ["compute_nodes.id"], ondelete="SET NULL" ), sa.ForeignKeyConstraint( ["training_job_id"], ["training_jobs.id"], ondelete="SET NULL" ), sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_fine_tuned_models_deleted_at"), "fine_tuned_models", ["deleted_at"], unique=False, ) create_index_if_not_exists( op.f("ix_fine_tuned_models_id"), "fine_tuned_models", ["id"], unique=False ) create_index_if_not_exists( op.f("ix_fine_tuned_models_training_job_id"), "fine_tuned_models", ["training_job_id"], unique=False, ) create_index_if_not_exists( op.f("ix_fine_tuned_models_user_id"), "fine_tuned_models", ["user_id"], unique=False, ) # ------------------------------------------------------------------------- # gpu_allocations # ------------------------------------------------------------------------- if not has_table("gpu_allocations"): op.create_table( "gpu_allocations", sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), sa.Column("node_id", sa.String(length=64), nullable=True), sa.Column("job_id", sa.String(length=64), nullable=True), sa.Column("model_id", sa.String(length=64), nullable=True), sa.Column("vram_reserved_gb", sa.Float(), nullable=False), sa.Column("created_at", sa.BigInteger(), nullable=True), sa.ForeignKeyConstraint( ["job_id"], ["training_jobs.id"], ondelete="CASCADE" ), sa.ForeignKeyConstraint( ["model_id"], ["fine_tuned_models.id"], ondelete="CASCADE" ), sa.ForeignKeyConstraint( ["node_id"], ["compute_nodes.id"], ondelete="CASCADE" ), sa.PrimaryKeyConstraint("id"), ) # ------------------------------------------------------------------------- # inference_deployments # ------------------------------------------------------------------------- if not has_table("inference_deployments"): op.create_table( "inference_deployments", sa.Column("id", sa.String(length=64), nullable=False), sa.Column("node_id", sa.String(length=128), nullable=True), sa.Column("internal_hostname", sa.String(length=128), nullable=True), sa.Column("base_model_id", sa.String(length=128), nullable=True), sa.Column("fine_tuned_model_id", sa.String(length=64), nullable=True), sa.Column("port", sa.Integer(), nullable=True), sa.Column("status", _STATUS_ENUM, nullable=True), sa.Column("current_throughput", sa.Float(), nullable=True), sa.Column("last_seen", sa.BigInteger(), nullable=True), sa.Column("tensor_parallel_size", sa.Integer(), nullable=False), sa.ForeignKeyConstraint(["base_model_id"], ["base_models.id"]), sa.ForeignKeyConstraint(["fine_tuned_model_id"], ["fine_tuned_models.id"]), sa.PrimaryKeyConstraint("id"), ) create_index_if_not_exists( op.f("ix_inference_deployments_id"), "inference_deployments", ["id"], unique=False, ) def downgrade() -> None: """Downgrade schema.""" drop_index_if_exists("ix_inference_deployments_id", "inference_deployments") drop_table_if_exists("inference_deployments") drop_table_if_exists("gpu_allocations") drop_index_if_exists("ix_fine_tuned_models_user_id", "fine_tuned_models") drop_index_if_exists("ix_fine_tuned_models_training_job_id", "fine_tuned_models") drop_index_if_exists("ix_fine_tuned_models_id", "fine_tuned_models") drop_index_if_exists("ix_fine_tuned_models_deleted_at", "fine_tuned_models") drop_table_if_exists("fine_tuned_models") drop_index_if_exists("ix_vector_store_files_id", "vector_store_files") drop_table_if_exists("vector_store_files") drop_table_if_exists("user_assistants") drop_index_if_exists("ix_training_jobs_user_id", "training_jobs") drop_index_if_exists("ix_training_jobs_id", "training_jobs") drop_index_if_exists("ix_training_jobs_deleted_at", "training_jobs") drop_index_if_exists("ix_training_jobs_dataset_id", "training_jobs") drop_index_if_exists("idx_trainingjob_user_id", "training_jobs") drop_index_if_exists("idx_trainingjob_status", "training_jobs") drop_table_if_exists("training_jobs") drop_table_if_exists("thread_participants") drop_index_if_exists("idx_file_storage_file_id", "file_storage") drop_table_if_exists("file_storage") drop_index_if_exists("ix_actions_tool_call_id", "actions") drop_index_if_exists("ix_actions_id", "actions") drop_index_if_exists("ix_actions_confidence_score", "actions") drop_table_if_exists("actions") drop_index_if_exists("ix_vector_stores_id", "vector_stores") drop_index_if_exists("ix_vector_stores_deleted_at", "vector_stores") drop_table_if_exists("vector_stores") drop_index_if_exists("ix_threads_owner_id", "threads") drop_index_if_exists("ix_threads_id", "threads") drop_table_if_exists("threads") drop_index_if_exists("ix_sandboxes_id", "sandboxes") drop_table_if_exists("sandboxes") drop_index_if_exists("ix_runs_user_id", "runs") drop_table_if_exists("runs") drop_index_if_exists("ix_files_id", "files") drop_index_if_exists("ix_files_deleted_at", "files") drop_table_if_exists("files") drop_index_if_exists("ix_datasets_user_id", "datasets") drop_index_if_exists("ix_datasets_id", "datasets") drop_index_if_exists("ix_datasets_file_id", "datasets") drop_index_if_exists("ix_datasets_deleted_at", "datasets") drop_table_if_exists("datasets") drop_index_if_exists("ix_batfish_snapshots_user_id", "batfish_snapshots") drop_index_if_exists("ix_batfish_snapshots_snapshot_key", "batfish_snapshots") drop_index_if_exists("ix_batfish_snapshots_id", "batfish_snapshots") drop_index_if_exists("idx_batfish_user_id", "batfish_snapshots") drop_index_if_exists("idx_batfish_status", "batfish_snapshots") drop_table_if_exists("batfish_snapshots") drop_index_if_exists("ix_audit_logs_user_id", "audit_logs") drop_index_if_exists("ix_audit_logs_timestamp", "audit_logs") drop_index_if_exists("ix_audit_logs_id", "audit_logs") drop_index_if_exists("ix_audit_logs_entity_type", "audit_logs") drop_index_if_exists("ix_audit_logs_entity_id", "audit_logs") drop_index_if_exists("ix_audit_logs_action", "audit_logs") drop_table_if_exists("audit_logs") drop_index_if_exists("ix_assistants_owner_id", "assistants") drop_index_if_exists("ix_assistants_id", "assistants") drop_table_if_exists("assistants") drop_index_if_exists("ix_api_keys_user_id", "api_keys") drop_index_if_exists("ix_api_keys_id", "api_keys") drop_index_if_exists("ix_api_keys_hashed_key", "api_keys") drop_index_if_exists("idx_apikey_user_id_active", "api_keys") drop_table_if_exists("api_keys") drop_index_if_exists("ix_users_provider_user_id", "users") drop_index_if_exists("ix_users_oauth_provider", "users") drop_index_if_exists("ix_users_id", "users") drop_index_if_exists("ix_users_email", "users") drop_index_if_exists("idx_user_is_admin", "users") drop_index_if_exists("idx_user_email", "users") drop_table_if_exists("users") drop_index_if_exists("ix_messages_id", "messages") drop_index_if_exists("ix_messages_assistant_id", "messages") drop_table_if_exists("messages") drop_index_if_exists("ix_compute_nodes_id", "compute_nodes") drop_table_if_exists("compute_nodes") drop_index_if_exists("ix_base_models_id", "base_models") drop_table_if_exists("base_models")