src/db/postgres/vector_store.py

"""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