diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py index 0e460f335..0ec114a8f 100644 --- a/src/langbot/pkg/persistence/tenant_uow.py +++ b/src/langbot/pkg/persistence/tenant_uow.py @@ -13,7 +13,7 @@ import typing import sqlalchemy import sqlalchemy.ext.asyncio as sqlalchemy_asyncio import sqlalchemy.orm as sqlalchemy_orm -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import HALFVEC, Vector from sqlalchemy.dialects.postgresql.dml import OnConflictDoNothing as PostgreSQLOnConflictDoNothing from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLOnConflictDoUpdate from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing @@ -281,7 +281,7 @@ def _validate_scoped_sql_type( return seen.add(identity) - if type(sql_type) is Vector: + if type(sql_type) in {Vector, HALFVEC}: return if not type(sql_type).__module__.startswith('sqlalchemy.'): raise ScopedSessionTransactionError('TenantUnitOfWork does not allow custom SQL types in public statements') @@ -462,7 +462,7 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s if isinstance(element, sqlalchemy.sql.elements.BindParameter) and element.literal_execute: raise ScopedSessionTransactionError('TenantUnitOfWork does not allow literal-execute SQL parameters') - if isinstance(element, sqlalchemy.sql.elements.Cast) and type(element.type) is not Vector: + if isinstance(element, sqlalchemy.sql.elements.Cast) and type(element.type) not in {Vector, HALFVEC}: raise ScopedSessionTransactionError( 'TenantUnitOfWork only allows the trusted pgvector cast used by tenant vector search' ) diff --git a/tests/unit_tests/persistence/test_tenant_uow.py b/tests/unit_tests/persistence/test_tenant_uow.py index bc8369f08..34e907095 100644 --- a/tests/unit_tests/persistence/test_tenant_uow.py +++ b/tests/unit_tests/persistence/test_tenant_uow.py @@ -7,7 +7,7 @@ from types import SimpleNamespace import pytest import sqlalchemy as sa -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import HALFVEC, Vector from sqlalchemy.dialects.postgresql import insert as postgresql_insert from sqlalchemy.dialects.sqlite import insert as sqlite_insert from sqlalchemy.ext.asyncio import create_async_engine @@ -967,6 +967,7 @@ async def test_scoped_session_rejects_raw_or_unapproved_sql( ), sa.select(sa.column('embedding').op('<=>')(sa.literal([0.1]))), sa.select(sa.cast(sa.column('embedding'), Vector(384))), + sa.select(sa.cast(sa.column('embedding'), HALFVEC(3072))), sa.insert(sa.table('rows', sa.column('id'))).values(id=1), _multi_value_statement(value=1), _on_conflict_statement(update_value=sa.func.coalesce(sa.literal(1), sa.literal(0))),