fix: allow trusted halfvec tenant search casts (#2402)

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

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

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