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

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

pyproject.toml

@@ -71,7 +71,7 @@ dependencies = [
     "chromadb>=1.0.0,<2.0.0",
     "qdrant-client (>=1.15.1,<2.0.0)",
     "pyseekdb==1.1.0.post3",
-    "langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@077a4185f26d19a505bacafd8fb86d6734b6344b",
+    "langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@9d216208cdfb41f0cb7fcb64632e2a46816d6dc6",
     "asyncpg>=0.30.0",
     "line-bot-sdk>=3.19.0",
     "matrix-nio>=0.25.2",

src/langbot/pkg/cloud/bootstrap.py

@@ -18,7 +18,7 @@ from .model_catalog import CloudModelCatalogProvider
 
 CLOUD_BOOTSTRAP_ENTRY_POINT = 'langbot.cloud_bootstrap'
 REQUIRED_TENANT_ISOLATION_VERSION = 2
-SUPPORTED_PGVECTOR_DIMENSIONS = frozenset({384, 512, 768, 1024, 1536})
+SUPPORTED_PGVECTOR_DIMENSIONS = frozenset({384, 512, 768, 1024, 1536, 3072})
 
 
 class CloudBootstrapError(RuntimeError):

src/langbot/pkg/persistence/alembic/versions/001a_pgvector_dimension_3072.py

@@ -0,0 +1,43 @@
+"""enable 3072-dimensional pgvector embeddings
+
+Revision ID: 001a_pgvector_dimension_3072
+Revises: 0019_single_workspace_owner
+Create Date: 2026-08-05
+"""
+
+from __future__ import annotations
+import sqlalchemy as sa
+from alembic import op
+
+revision = '001a_pgvector_dimension_3072'
+down_revision = '0019_single_workspace_owner'
+branch_labels = None
+depends_on = None
+_TABLE = 'langbot_vectors'
+_CHECK = 'ck_langbot_vectors_embedding_dimension_enabled'
+_INDEX = 'ix_langbot_vectors_hnsw_cosine_3072'
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    if conn.dialect.name != 'postgresql' or _TABLE not in sa.inspect(conn).get_table_names():
+        return
+    op.drop_constraint(_CHECK, _TABLE, type_='check')
+    op.create_check_constraint(_CHECK, _TABLE, 'embedding_dimension IN (384, 512, 768, 1024, 1536, 3072)')
+    op.execute(
+        sa.text(
+            f'CREATE INDEX {_INDEX} ON {_TABLE} USING hnsw ((embedding::halfvec(3072)) halfvec_cosine_ops) WHERE embedding_dimension = 3072'
+        )
+    )
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    if conn.dialect.name != 'postgresql' or _TABLE not in sa.inspect(conn).get_table_names():
+        return
+    count = conn.scalar(sa.text(f'SELECT COUNT(*) FROM {_TABLE} WHERE embedding_dimension = 3072'))
+    if count:
+        raise RuntimeError('Cannot disable 3072-dimensional pgvector while matching embeddings exist')
+    op.drop_index(_INDEX, table_name=_TABLE)
+    op.drop_constraint(_CHECK, _TABLE, type_='check')
+    op.create_check_constraint(_CHECK, _TABLE, 'embedding_dimension IN (384, 512, 768, 1024, 1536)')

src/langbot/pkg/persistence/mgr.py

@@ -98,7 +98,7 @@ _WORKSPACE_ALEMBIC_REVISION = '0009_workspace_tenancy'
 _RESOURCE_SCOPE_ALEMBIC_REVISION = '0010_scope_resources'
 _OSS_WORKSPACE_METADATA_KEY = 'oss_workspace_uuid'
 _RELEASE_MIGRATION_ADVISORY_LOCK_ID = 0x4C414E47424F5432
-_PGVECTOR_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536)
+_PGVECTOR_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536, 3072)
 _RUNTIME_SCHEMA = 'public'
 _ALEMBIC_RUNTIME_TABLE = 'alembic_version'
 _RUNTIME_TABLE_PRIVILEGES = frozenset({'SELECT', 'INSERT', 'UPDATE', 'DELETE'})
@@ -1356,14 +1356,16 @@ class PersistenceManager:
             index = by_index.get(index_name)
             index_definition = normalized(None if index is None else index['definition'])
             predicate = normalized(None if index is None else index['predicate'])
+            vector_type = 'halfvec' if dimension > 2000 else 'vector'
+            operator_class = f'{vector_type}_cosine_ops'
             if (
                 index is None
                 or index['access_method'] != 'hnsw'
                 or index['is_valid'] is not True
                 or index['is_ready'] is not True
-                or f'vector({dimension})' not in index_definition
-                or f'(embedding)::vector({dimension})' not in index_definition
-                or 'vector_cosine_ops' not in index_definition
+                or f'{vector_type}({dimension})' not in index_definition
+                or f'(embedding)::{vector_type}({dimension})' not in index_definition
+                or operator_class not in index_definition
                 or predicate.strip('() ') != f'embedding_dimension = {dimension}'
             ):
                 raise RuntimeError(f'PostgreSQL pgvector ANN index {index_name!r} is invalid')

src/langbot/pkg/vector/mgr.py

@@ -67,7 +67,7 @@ class VectorDBManager:
                 use_business_database = pgvector_config.get('use_business_database', False)
                 allowed_dimensions = pgvector_config.get(
                     'allowed_dimensions',
-                    [384, 512, 768, 1024, 1536],
+                    [384, 512, 768, 1024, 1536, 3072],
                 )
                 common_options = {
                     'use_business_database': use_business_database,

src/langbot/pkg/vector/vdbs/pgvector_db.py

@@ -6,7 +6,7 @@ from collections.abc import AsyncIterator
 from typing import Any
 
 import sqlalchemy
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import HALFVEC, Vector
 from sqlalchemy.dialects.postgresql import insert as postgresql_insert
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
 from sqlalchemy.orm import declarative_base
@@ -18,7 +18,7 @@ from langbot.pkg.vector.vdb import VectorDatabase
 
 Base = declarative_base()
 
-DEFAULT_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536)
+DEFAULT_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536, 3072)
 
 # pgvector schema only stores these metadata fields.
 _PG_SUPPORTED_FIELDS = {'text', 'file_id', 'chunk_uuid'}
@@ -321,7 +321,12 @@ class PgVectorDatabase(VectorDatabase):
         if len(query_embedding) != scope.embedding_dimension:
             raise ValueError(f'Query embedding must have the selected dimension {scope.embedding_dimension}')
 
-        typed_embedding = sqlalchemy.cast(PgVectorEntry.embedding, Vector(scope.embedding_dimension))
+        typed_embedding = sqlalchemy.cast(
+            PgVectorEntry.embedding,
+            HALFVEC(scope.embedding_dimension)
+            if scope.embedding_dimension > 2000
+            else Vector(scope.embedding_dimension),
+        )
         distance = typed_embedding.cosine_distance(query_embedding)
         statement = (
             sqlalchemy.select(

src/langbot/templates/config.yaml

@@ -201,7 +201,7 @@ vdb:
         # keep this false when deliberately using an external pgvector DB.
         use_business_database: false
         # Release migrations create one partial ANN index per enabled value.
-        allowed_dimensions: [384, 512, 768, 1024, 1536]
+        allowed_dimensions: [384, 512, 768, 1024, 1536, 3072]
         host: '127.0.0.1'
         port: 5433
         database: 'langbot'

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

uv.lock

@@ -2125,7 +2125,7 @@ requires-dist = [
     { name = "ebooklib", specifier = ">=0.18" },
     { name = "gewechat-client", specifier = ">=0.1.5" },
     { name = "html2text", specifier = ">=2024.2.26" },
-    { name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=077a4185f26d19a505bacafd8fb86d6734b6344b" },
+    { name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=9d216208cdfb41f0cb7fcb64632e2a46816d6dc6" },
     { name = "langchain", specifier = ">=1.3.9" },
     { name = "langchain-core", specifier = ">=1.3.3" },
     { name = "langchain-text-splitters", specifier = ">=1.1.2" },
@@ -2192,7 +2192,7 @@ dev = [
 [[package]]
 name = "langbot-plugin"
 version = "0.5.0"
-source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=077a4185f26d19a505bacafd8fb86d6734b6344b#077a4185f26d19a505bacafd8fb86d6734b6344b" }
+source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=9d216208cdfb41f0cb7fcb64632e2a46816d6dc6#9d216208cdfb41f0cb7fcb64632e2a46816d6dc6" }
 dependencies = [
     { name = "aiofiles" },
     { name = "aiohttp" },