File size: 1,137 Bytes
bef5e76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | """PGVector store setup for document embeddings."""
from langchain_postgres import PGVector
from langchain_openai import AzureOpenAIEmbeddings
from src.config.settings import settings
from src.db.postgres.connection import _pgvector_engine
# Initialize embeddings
embeddings = AzureOpenAIEmbeddings(
azure_deployment=settings.azureai_deployment_name_embedding,
openai_api_version=settings.azureai_api_version_embedding,
azure_endpoint=settings.azureai_endpoint_url_embedding,
api_key=settings.azureai_api_key_embedding
)
# Use psycopg3 connection string (not asyncpg engine) with async_mode=True.
# psycopg3 supports multi-statement SQL, which PGVector needs for
# advisory_lock + CREATE EXTENSION vector. asyncpg rejects this as a prepared statement.
vector_store = PGVector(
embeddings=embeddings,
connection=_pgvector_engine,
collection_name="document_embeddings",
use_jsonb=True,
async_mode=True,
create_extension=False, # Extension pre-created in init_db.py (avoids multi-statement asyncpg bug)
)
def get_vector_store():
"""Get the vector store instance."""
return vector_store
|