pgvector_connection.py

from langchain.vectorstores.pgvector import PGVector, DistanceStrategy
from langchain.embeddings import HuggingFaceEmbeddings

model = 'bert-base-uncased'
embedding_model = HuggingFaceEmbeddings(model_name=model)

COLLECTION_NAME = 'top_10_dataset'
CONNECTION_STRING = PGVector.connection_string_from_db_params(
    driver='psycopg2',
    host=os.environ.get('PGVECTOR_HOST'),
    port=int(os.environ.get('PGVECTOR_PORT')),
    database=os.environ.get('PGVECTOR_DATABASE'),
    user=os.environ.get('PGVECTOR_USER'),
    password=os.environ.get('PGVECTOR_PASSWORD'),
)

pgvector = PGVector(
    collection_name=COLLECTION_NAME,
    embedding=embedding_model,
    connection_string=CONNECTION_STRING,
    distance_strategy=DistanceStrategy.COSINE
    )