migrations/versions/db67202be996_baseline.py
Rama main · réplica de lectura del repositorio original.
"""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")