| 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() |
