postgres

langroid/vector_store/postgres.py

PostgresDB(config=PostgresDBConfig())

Bases: VectorStore

Source code in langroid/vector_store/postgres.py

def __init__(self, config: PostgresDBConfig = PostgresDBConfig()):
    super().__init__(config)
    if not has_postgres:
        raise LangroidImportError("pgvector", "postgres")
    self.config: PostgresDBConfig = config
    self.engine = self._create_engine()
    PostgresDB._create_vector_extension(self.engine)
    self.SessionLocal = sessionmaker(
        autocommit=False, autoflush=False, bind=self.engine
    )
    self.metadata = MetaData()
    self._setup_table()

index_exists(connection, index_name)

Check if an index exists.

Source code in langroid/vector_store/postgres.py

def index_exists(self, connection: Connection, index_name: str) -> bool:
    """Check if an index exists."""
    query = text(
        "SELECT 1 FROM pg_indexes WHERE indexname = :index_name"
    ).bindparams(index_name=index_name)
    result = connection.execute(query).scalar()
    return bool(result)

delete_collection(collection_name)

Deletes a collection and its associated HNSW index, handling metadata synchronization issues.

Source code in langroid/vector_store/postgres.py

def delete_collection(self, collection_name: str) -> None:
    """
    Deletes a collection and its associated HNSW index, handling metadata
    synchronization issues.
    """
    with self.engine.connect() as connection:
        connection.execute(text("COMMIT"))
        index_name = f"hnsw_index_{collection_name}_embedding"
        drop_index_query = text(f"DROP INDEX CONCURRENTLY IF EXISTS {index_name}")
        connection.execute(drop_index_query)

        # 3. Now, drop the table using SQLAlchemy
        table = Table(collection_name, self.metadata)
        table.drop(self.engine, checkfirst=True)

        # 4. Refresh metadata again after dropping the table
        self.metadata.clear()
        self.metadata.reflect(bind=self.engine)