flask-migrate 未导入所需的库 pgvector

Question

我在尝试使用带有flask和flask-sqlalchemy的pgvector创建向量嵌入时收到pgvector的NameError。

NameError: name 'pgvector' is not defined

#extensions.py

from flask_sqlalchemy import SQLAlchemy
db = SQLAlchemy()

#models.py

from extensions import db
from sqlalchemy.orm import mapped_column
from pgvector.sqlalchemy import Vector

N_DIM = 1536

class InstagramPosts(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    embedding = mapped_column(Vector(N_DIM))

然后我通过“flask db migrate -m”added embedding”迁移到我的 postgres 数据库 并尝试通过“flask db升级”运行升级

之后我收到名称错误。

[堆栈跟踪]

"~/migrations/versions/2b0a38ba6267_added_pgvector_column.py", line 23, in upgrade
    batch_op.add_column(sa.Column('embedding', pgvector.sqlalchemy.vector.VECTOR(dim=1536), nullable=True))
                                               ^^^^^^^^
NameError: name 'pgvector' is not defined

我已经通过

在 postgres 数据库中创建了扩展

CREATE EXTENSION IF NOT EXISTS vector;

我已经通过 pip install pgvector 安装了 pgvector，并且通过 pip freeze 在我安装的库中可以看到它

不确定它是否相关，但就其价值而言，我的应用程序是带有蓝图的 Flask 应用程序，我正在使用应用程序工厂方法来创建应用程序等。

编辑----------------

经过进一步调查，pgvector 似乎没有按应有的方式自动导入到迁移文件中。我在下面添加了迁移文件。如果我手动打开迁移文件并添加“导入 pgvector”，它就可以工作。所以我现在的问题是为什么 pgvector 没有像应该那样用 Flask-migrate 导入。

#迁移文件85a0af23e160_added_embeddings_column.py

"""added embeddings column

Revision ID: 85a0af23e160
Revises: f768e743870b
Create Date: 2024-08-15 19:39:52.497427

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '85a0af23e160'
down_revision = 'f768e743870b'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('instagram_posts', schema=None) as batch_op:
        batch_op.add_column(sa.Column('embedding', pgvector.sqlalchemy.vector.VECTOR(dim=1536), nullable=True))

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('instagram_posts', schema=None) as batch_op:
        batch_op.drop_column('embedding')

    # ### end Alembic commands ###

Answer 1

好的。终于明白了这一点。 解决方案：编辑迁移 script.py.mako 文件以包含任何必要的导入。

#script.py.mako

"""${message}

Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}

"""
from alembic import op
import sqlalchemy as sa
import pgvector     #  <---- This line was added
${imports if imports else ""}

# revision identifiers, used by Alembic.
revision = ${repr(up_revision)}
down_revision = ${repr(down_revision)}
branch_labels = ${repr(branch_labels)}
depends_on = ${repr(depends_on)}


def upgrade():
    ${upgrades if upgrades else "pass"}


def downgrade():
    ${downgrades if downgrades else "pass"}