Added literal binds support for SQLAlchemy - closes #51

Author: ankane

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.2.5 (unreleased)
+
+- Added literal binds support for SQLAlchemy
+
 ## 0.2.4 (2023-11-24)
 
 - Improved reflection with SQLAlchemy

pgvector/sqlalchemy/__init__.py

@@ -1,12 +1,13 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
-from sqlalchemy.types import UserDefinedType, Float
+from sqlalchemy.types import UserDefinedType, Float, String
 from ..utils import from_db, to_db
 
 __all__ = ['Vector']
 
 
 class Vector(UserDefinedType):
     cache_ok = True
+    _string = String()
 
     def __init__(self, dim=None):
         super(UserDefinedType, self).__init__()
@@ -22,6 +23,12 @@ def process(value):
             return to_db(value, self.dim)
         return process
 
+    def literal_processor(self, dialect):
+        string_literal_processor = self._string._cached_literal_processor(dialect)
+        def process(value):
+            return string_literal_processor(to_db(value, self.dim))
+        return process
+
     def result_processor(self, dialect, coltype):
         def process(value):
             return from_db(value)

tests/test_sqlalchemy.py

@@ -219,6 +219,10 @@ def test_inspect(self):
         columns = inspect(engine).get_columns('orm_item')
         assert isinstance(columns[1]['type'], Vector)
 
+    def test_literal_binds(self):
+        sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(compile_kwargs={'literal_binds': True})
+        assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql)
+
     @pytest.mark.asyncio
     async def test_async(self):
         engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test')