feat(backend): add pgvector-python

aigis-llm · Jan 28, 2025 · 17272e7 · 17272e7
1 parent 3b63851
commit 17272e7
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -7,6 +7,7 @@ authors = [
 ]
 dependencies = [
     "psycopg[binary]>=3.2.3",
+    "pgvector>=0.3.6",
     "numpy>=2.1.3",
     "chonkie[semantic]>=0.4.0",
     "fastapi[standard]>=0.115.5",

diff --git a/src/aigis/utils.py b/src/aigis/utils.py
@@ -5,6 +5,9 @@
 
 from fastapi import Depends, Request
 from openai import AsyncOpenAI
+from pgvector.psycopg import (  # pyright: ignore [reportMissingTypeStubs]
+	register_vector_async,  # pyright: ignore [reportUnknownVariableType]
+)
 from psycopg import AsyncConnection
 
 from supabase import AsyncClient, AsyncClientOptions, create_async_client
@@ -59,6 +62,7 @@ async def openai_client(_request: Request) -> AsyncGenerator[AsyncOpenAI]:
 @asynccontextmanager
 async def postgres_client(_request: Request) -> AsyncGenerator[AsyncConnection]:
 	conn = await AsyncConnection.connect(postgres_url)
+	await register_vector_async(conn)
 	try:
 		yield conn
 	finally:

diff --git a/uv.lock b/uv.lock