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

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

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(