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],
                     },
                 },
             }

tests/unit_tests/cloud/test_bootstrap.py

@@ -108,7 +108,7 @@ def _cloud_config() -> dict:
             'use': 'pgvector',
             'pgvector': {
                 'use_business_database': True,
-                'allowed_dimensions': [384, 768, 1536],
+                'allowed_dimensions': [384, 768, 1536, 3072],
             },
         },
         'mcp': {'stdio': {'enabled': False}},
@@ -216,7 +216,6 @@ async def test_cloud_directory_capacity_contract_is_fail_closed(directory_config
     [
         ({'use_business_database': False, 'allowed_dimensions': [1536]}, 'use_business_database=true'),
         ({'use_business_database': True, 'allowed_dimensions': []}, 'allowed_dimensions'),
-        ({'use_business_database': True, 'allowed_dimensions': [3072]}, 'allowed_dimensions'),
         ({'use_business_database': True, 'allowed_dimensions': [True]}, 'allowed_dimensions'),
     ],
 )

tests/unit_tests/vector/test_mgr.py

@@ -213,7 +213,7 @@ class TestVectorDBManagerInitialization:
                 mock_app,
                 connection_string='postgresql://user:pass@host:5432/langbot',
                 use_business_database=False,
-                allowed_dimensions=[384, 512, 768, 1024, 1536],
+                allowed_dimensions=[384, 512, 768, 1024, 1536, 3072],
             )
 
     def test_initialize_pgvector_with_individual_params(self):
@@ -251,7 +251,7 @@ class TestVectorDBManagerInitialization:
                 user='admin',
                 password='secret',
                 use_business_database=False,
-                allowed_dimensions=[384, 512, 768, 1024, 1536],
+                allowed_dimensions=[384, 512, 768, 1024, 1536, 3072],
             )
 
     def test_initialize_pgvector_defaults(self):
@@ -280,7 +280,7 @@ class TestVectorDBManagerInitialization:
                 user='postgres',
                 password='postgres',
                 use_business_database=False,
-                allowed_dimensions=[384, 512, 768, 1024, 1536],
+                allowed_dimensions=[384, 512, 768, 1024, 1536, 3072],
             )
 
     def test_initialize_pgvector_with_shared_business_database(self):