open-webui/backend/open_webui/migrations/versions/6a39f3d8e55c_add_knowledge_table.py
2024-10-04 00:59:19 -07:00

81 lines
2.4 KiB
Python

"""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")