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

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

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