backend/open_webui/migrations/versions/6a39f3d8e55c_add_knowledge_table.py

"""Add knowledge table

Revision ID: 6a39f3d8e55c
Revises: c0fbf31ca0db
Create Date: 2024-10-01 14:02:35.241684

"""

from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import table, column, select
import json

revision = "6a39f3d8e55c"
down_revision = "c0fbf31ca0db"
branch_labels = None
depends_on = None


def upgrade():
    # Creating the 'knowledge' table
    print("Creating knowledge table")
    knowledge_table = op.create_table(
        "knowledge",
        sa.Column("id", sa.Text(), primary_key=True),
        sa.Column("user_id", sa.Text(), nullable=False),
        sa.Column("name", sa.Text(), nullable=False),
        sa.Column("description", sa.Text(), nullable=True),
        sa.Column("data", sa.JSON(), nullable=True),
        sa.Column("meta", sa.JSON(), nullable=True),
        sa.Column("created_at", sa.BigInteger(), nullable=False),
        sa.Column("updated_at", sa.BigInteger(), nullable=True),
    )

    print("Migrating data from document table to knowledge table")
    # Representation of the existing 'document' table
    document_table = table(
        "document",
        column("collection_name", sa.String()),
        column("user_id", sa.String()),
        column("name", sa.String()),
        column("title", sa.Text()),
        column("content", sa.Text()),
        column("timestamp", sa.BigInteger()),
    )

    # Select all from existing document table
    documents = op.get_bind().execute(
        select(
            document_table.c.collection_name,
            document_table.c.user_id,
            document_table.c.name,
            document_table.c.title,
            document_table.c.content,
            document_table.c.timestamp,
        )
    )

    # Insert data into knowledge table from document table
    for doc in documents:
        op.get_bind().execute(
            knowledge_table.insert().values(
                id=doc.collection_name,
                user_id=doc.user_id,
                description=doc.name,
                meta={
                    "legacy": True,
                    "document": True,
                    "tags": json.loads(doc.content or "{}").get("tags", []),
                },
                name=doc.title,
                created_at=doc.timestamp,
                updated_at=doc.timestamp,  # using created_at for both created_at and updated_at in project
            )
        )


def downgrade():
    op.drop_table("knowledge")