fix: support 3072-dimensional knowledge embeddings (#2401)

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>

tests/integration/persistence/test_migrations.py

@@ -105,7 +105,7 @@ class TestSQLiteMigrationUpgrade:
         await run_alembic_upgrade(sqlite_engine, 'head')
 
         assert await get_alembic_current(sqlite_engine) == _get_script_head()
-        assert _get_script_head() == '0019_single_workspace_owner'
+        assert _get_script_head() == '001a_pgvector_dimension_3072'
 
     @pytest.mark.asyncio
     async def test_upgrade_from_baseline_to_head(self, sqlite_engine):

tests/integration/persistence/test_pgvector_postgres.py

@@ -85,6 +85,32 @@ async def clean_database(postgres_engine: AsyncEngine):
     await clean()
 
 
+async def test_upgrade_adds_3072_dimension_index_and_constraint(
+    postgres_engine: AsyncEngine,
+    clean_database,
+) -> None:
+    async with postgres_engine.begin() as conn:
+        await conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
+        await conn.run_sync(Base.metadata.create_all)
+    await run_alembic_stamp(postgres_engine, '0010_scope_resources')
+    await run_alembic_upgrade(postgres_engine, 'head')
+
+    async with postgres_engine.connect() as conn:
+        constraint = await conn.scalar(
+            text(
+                'SELECT pg_get_constraintdef(oid) FROM pg_constraint '
+                "WHERE conrelid = 'langbot_vectors'::regclass "
+                "AND conname = 'ck_langbot_vectors_embedding_dimension_enabled'"
+            )
+        )
+        assert '3072' in constraint
+        index_definition = await conn.scalar(
+            text("SELECT indexdef FROM pg_indexes WHERE indexname = 'ix_langbot_vectors_hnsw_cosine_3072'")
+        )
+        assert 'halfvec(3072)' in index_definition
+        assert 'halfvec_cosine_ops' in index_definition
+
+
 async def test_legacy_upgrade_temporarily_suspends_and_restores_source_rls_for_unprivileged_owner(
     postgres_url: str,
     postgres_engine: AsyncEngine,

tests/integration/persistence/test_release_migration_postgres.py

@@ -92,7 +92,7 @@ def _application(postgres_url: str, *, runtime_role: str = 'langbot_runtime_not_
                     'use': 'pgvector',
                     'pgvector': {
                         'use_business_database': True,
-                        'allowed_dimensions': [384, 512, 768, 1024, 1536],
+                        'allowed_dimensions': [384, 512, 768, 1024, 1536, 3072],
                     },
                 },
             }