Update pgvector

Author: pamelafox

src/backend/fastapi_app/postgres_engine.py

@@ -2,7 +2,9 @@
 import os
 
 from azure.identity import AzureDeveloperCliCredential
+from pgvector.asyncpg import register_vector
 from sqlalchemy import event
+from sqlalchemy.engine import AdaptedConnection
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
 from fastapi_app.dependencies import get_azure_credential
@@ -35,6 +37,10 @@ def get_password_from_azure_credential():
         echo=False,
     )
 
+    @event.listens_for(engine.sync_engine, "connect")
+    def register_custom_types(dbapi_connection: AdaptedConnection, *args):
+        dbapi_connection.run_async(register_vector)
+
     @event.listens_for(engine.sync_engine, "do_connect")
     def update_password_token(dialect, conn_rec, cargs, cparams):
         if token_based_password:

src/backend/fastapi_app/postgres_searcher.py

@@ -1,5 +1,5 @@
+import numpy as np
 from openai import AsyncAzureOpenAI, AsyncOpenAI
-from pgvector.utils import to_db
 from sqlalchemy import Float, Integer, column, select, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -87,7 +87,7 @@ async def search(
         results = (
             await self.db_session.execute(
                 sql,
-                {"embedding": to_db(query_vector), "query": query_text, "k": 60},
+                {"embedding": np.array(query_vector), "query": query_text, "k": 60},
             )
         ).fetchall()
 

src/backend/pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "aiohttp>=3.9.5,<4.0.0",
     "asyncpg>=0.29.0,<1.0.0",
     "SQLAlchemy[asyncio]>=2.0.30,<3.0.0",
-    "pgvector>=0.2.5,<0.3.0",
+    "pgvector>=0.3.0,<0.4.0",
     "openai>=1.34.0,<2.0.0",
     "tiktoken>=0.7.0,<0.8.0",
     "openai-messages-token-helper>=0.1.8,<0.2.0",