cack / db.py
ujalaarshad17's picture
Upload 8 files
a16eb78 verified
from langchain_community.vectorstores.pgvector import PGVector
import os
# try:
# CONNECTION_STRING = PGVector.connection_string_from_db_params(
# driver="psycopg2",
# host=os.getenv("POSTGRES_HOST"),
# port=int(os.getenv("POSTGRES_PORT", 5432)),
# database=os.getenv("POSTGRES_DB"),
# user=os.getenv("POSTGRES_USER"),
# password=os.getenv("POSTGRES_PASSWORD"),
# )
# print("Successfully established the connection")
# except Exception as e:
# print("Error in establishing the connection with DB: {e}")
print("Entered here")
from dotenv import load_dotenv
load_dotenv()
from langchain_google_genai import GoogleGenerativeAIEmbeddings
embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
from sqlalchemy import create_engine
from langchain_postgres.vectorstores import PGVector
from langchain_core.documents import Document
document_1 = Document(page_content="fddsdfoo", metadata={"baz": "bar"})
document_2 = Document(page_content="thufeed", metadata={"bar": "baz"})
document_3 = Document(page_content="i wefsill be deleted :(")
documents = [document_1, document_2, document_3]
ids = ["1", "2", "3"]
engine = create_engine(os.environ['CONNECTION_STRING'])
vector_store = PGVector.from_documents(
documents=documents,
embedding=embeddings,
connection=os.environ['CONNECTION_STRING'],
collection_name="collection_name",
use_jsonb=True,
)
# vector_store.add_documents(documents=documents, ids=ids)
print("Stored babe")