File size: 1,315 Bytes
b88c922 | 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 33 34 35 36 37 38 | import sqlalchemy
from langchain.vectorstores.pgvector import PGVector as Base,CollectionStore
from sqlalchemy.orm import Session
class PGVector(Base):
def __init__(
self,
pre_delete_embeddings: bool = False,
**kwargs
) -> None:
self.pre_delete_embeddings = pre_delete_embeddings
Base.__init__(self, **kwargs)
def connect(self) -> sqlalchemy.engine.Connection:
engine = sqlalchemy.create_engine(self.connection_string, echo=True)
conn = engine.connect()
return conn
def delete_embeddings(self) -> None:
self.logger.debug("Trying to delete embeddings")
with Session(self._conn) as session:
collection = self.get_collection(session)
if not collection:
self.logger.error("Collection not found")
return
session.delete(collection.embeddings)
session.commit()
def create_collection(self) -> None:
if self.pre_delete_collection:
self.delete_collection()
with Session(self._conn) as session:
CollectionStore.get_or_create(
session, self.collection_name, cmetadata=self.collection_metadata
)
if self.pre_delete_embeddings:
self.delete_embeddings() |