Added OpenAI large embedding / improve migration and upgrade handling for containers

migrations/tenant/env.py

@@ -99,6 +99,16 @@ def run_migrations_online():
         poolclass=NullPool,
     )
 
+    def process_revision_directives(context, revision, directives):
+        if config.cmd_opts.autogenerate:
+            script = directives[0]
+            if script.upgrade_ops is not None:
+                # Add import for pgvector and set search path
+                script.upgrade_ops.ops.insert(0, sa.schema.ExecuteSQLOp(
+                    "SET search_path TO CURRENT_SCHEMA(), public; IMPORT pgvector",
+                    execution_options=None
+                ))
+
     with connectable.connect() as connection:
         print(tenants)
         for tenant in tenants:
@@ -106,17 +116,18 @@ def run_migrations_online():
             # set search path on the connection, which ensures that
             # PostgreSQL will emit all CREATE / ALTER / DROP statements
             # in terms of this schema by default
-            connection.execute(text(f'SET search_path TO "{tenant}"'))
+            connection.execute(text(f'SET search_path TO "{tenant}", public'))
             # in SQLAlchemy v2+ the search path change needs to be committed
             connection.commit()
 
             # make use of non-supported SQLAlchemy attribute to ensure
             # the dialect reflects tables in terms of the current tenant name
-            connection.dialect.default_schema_name = tenant
+            connection.dialect.default_schema_name = str(tenant)
 
             context.configure(
                 connection=connection,
                 target_metadata=get_metadata(),
+                process_revision_directives=process_revision_directives,
             )
 
             with context.begin_transaction():

migrations/tenant/script.py.mako

@@ -7,6 +7,7 @@ Create Date: ${create_date}
 """
 from alembic import op
 import sqlalchemy as sa
+import pgvector
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.

migrations/tenant/versions/5d5437d81041_add_embedding_model_for_openai_large_.py

@@ -0,0 +1,34 @@
+"""Add Embedding model for OpenAI large model
+
+Revision ID: 5d5437d81041
+Revises: d173cea8d204
+Create Date: 2024-07-04 15:49:43.665685
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import pgvector
+
+
+# revision identifiers, used by Alembic.
+revision = '5d5437d81041'
+down_revision = 'd173cea8d204'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('embedding_large_openai',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('embedding', pgvector.sqlalchemy.Vector(dim=3072), nullable=False),
+    sa.ForeignKeyConstraint(['id'], ['embeddings.id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('embedding_large_openai')
+    # ### end Alembic commands ###