From cdd5c6589c356d197cc4f92c3d8a338de5166a77 Mon Sep 17 00:00:00 2001 From: Hyu Date: Wed, 5 Aug 2026 21:19:57 +0800 Subject: [PATCH] fix: support 3072-dimensional knowledge embeddings (#2401) Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com> --- pyproject.toml | 2 +- src/langbot/pkg/cloud/bootstrap.py | 2 +- .../versions/001a_pgvector_dimension_3072.py | 43 +++++++++++++++++++ src/langbot/pkg/persistence/mgr.py | 10 +++-- src/langbot/pkg/vector/mgr.py | 2 +- src/langbot/pkg/vector/vdbs/pgvector_db.py | 11 +++-- src/langbot/templates/config.yaml | 2 +- .../persistence/test_migrations.py | 2 +- .../persistence/test_pgvector_postgres.py | 26 +++++++++++ .../test_release_migration_postgres.py | 2 +- tests/unit_tests/cloud/test_bootstrap.py | 3 +- tests/unit_tests/vector/test_mgr.py | 6 +-- uv.lock | 4 +- 13 files changed, 95 insertions(+), 20 deletions(-) create mode 100644 src/langbot/pkg/persistence/alembic/versions/001a_pgvector_dimension_3072.py diff --git a/pyproject.toml b/pyproject.toml index cd5f05449..b6144f3cb 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/langbot/pkg/cloud/bootstrap.py b/src/langbot/pkg/cloud/bootstrap.py index 2c7ddafd2..c631772ea 100644 --- a/src/langbot/pkg/cloud/bootstrap.py +++ b/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): diff --git a/src/langbot/pkg/persistence/alembic/versions/001a_pgvector_dimension_3072.py b/src/langbot/pkg/persistence/alembic/versions/001a_pgvector_dimension_3072.py new file mode 100644 index 000000000..7fc9c802a --- /dev/null +++ b/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)') diff --git a/src/langbot/pkg/persistence/mgr.py b/src/langbot/pkg/persistence/mgr.py index 121f19de5..e9062c066 100644 --- a/src/langbot/pkg/persistence/mgr.py +++ b/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') diff --git a/src/langbot/pkg/vector/mgr.py b/src/langbot/pkg/vector/mgr.py index 7a748d325..21e868223 100644 --- a/src/langbot/pkg/vector/mgr.py +++ b/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, diff --git a/src/langbot/pkg/vector/vdbs/pgvector_db.py b/src/langbot/pkg/vector/vdbs/pgvector_db.py index 341a99a96..b1c6fb3e4 100644 --- a/src/langbot/pkg/vector/vdbs/pgvector_db.py +++ b/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( diff --git a/src/langbot/templates/config.yaml b/src/langbot/templates/config.yaml index b0bb10c93..1ae2ba515 100644 --- a/src/langbot/templates/config.yaml +++ b/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' diff --git a/tests/integration/persistence/test_migrations.py b/tests/integration/persistence/test_migrations.py index 933d6fcd4..6e2b723ff 100644 --- a/tests/integration/persistence/test_migrations.py +++ b/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): diff --git a/tests/integration/persistence/test_pgvector_postgres.py b/tests/integration/persistence/test_pgvector_postgres.py index 7764fb7e0..1da6e4be0 100644 --- a/tests/integration/persistence/test_pgvector_postgres.py +++ b/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, diff --git a/tests/integration/persistence/test_release_migration_postgres.py b/tests/integration/persistence/test_release_migration_postgres.py index b6637afda..927f65127 100644 --- a/tests/integration/persistence/test_release_migration_postgres.py +++ b/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], }, }, } diff --git a/tests/unit_tests/cloud/test_bootstrap.py b/tests/unit_tests/cloud/test_bootstrap.py index 6225ba7b7..56fd4b9ca 100644 --- a/tests/unit_tests/cloud/test_bootstrap.py +++ b/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'), ], ) diff --git a/tests/unit_tests/vector/test_mgr.py b/tests/unit_tests/vector/test_mgr.py index 608f35121..84fcec96d 100644 --- a/tests/unit_tests/vector/test_mgr.py +++ b/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): diff --git a/uv.lock b/uv.lock index ece51234b..14eb1251a 100644 --- a/uv.lock +++ b/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" },