[PR #64] [MERGED] Add async mode for pgvector (v2) #142

New Issue

2026-02-16T05:16:44-05:00

yindo commented

2026-02-16 05:16:44 -05:00

📋 Pull Request Information

Original PR: https://github.com/langchain-ai/langchain-postgres/pull/64
Author: @pprados
Created: 6/10/2024
Status: ✅ Merged
Merged: 6/10/2024
Merged by: @eyurtsev

Base: main ← Head: pprados/async2

📝 Commits (10+)

188c83d Pre integration PPR
69c22a3 Add async mode
0abd730 Fix lint
e1ad8c8 It's possible to share the session_maker with sync mode.
dc6187a Fix apost_init
14d6da8 Add async API
e95df2f Rebase
9c2035c Fix create_vector_extension with async
7cc0bb1 Fix create_vector_extension with async
4cbeeb2 Re-ajuste les UI vis à vis de l'async (partiellement)

📊 Changes

3 files changed (+1328 additions, -86 deletions)

View changed files

📝 langchain_postgres/vectorstores.py (+830 -74)
📝 pyproject.toml (+4 -0)
📝 tests/unit_tests/test_vectorstore.py (+494 -12)

📄 Description

I made the mistake of doing a rebase while a review was in progress. This seems to block the process. There's a ‘1 change requested’ request that I can't validate. I propose another PR identical to the previous one.

This PR adds the async approach for pgvector.

Some remarks:

We use assert to check invocations and not if. Thus, in production, it is possible to
remove these checks with python -O ...
We propose a public session_maker attribute. This is very important for resilient
scenarios.

In a RAG architecture, it is necessary to import document chunks.

To keep track of the links between chunks and documents, we can use the
index() API.
This API proposes to use an SQL-type record manager.

In a classic use case, using SQLRecordManager and a vector database, it is impossible
to guarantee the consistency of the import.

Indeed, if a crash occurs during the import, there is an inconsistency between the SQL
database and the vector database.

PGVector is the solution to this problem.

Indeed, it is possible to use a single database (and not a two-phase commit with 2
technologies, if they are both compatible). But, for this, it is necessary to be able
to combine the transactions between the use of SQLRecordManager and PGVector as a
vector database.

This is only possible if it is possible to intervene on the session_maker.

This is why we propose to make this attribute public. By unifying the session_maker
of SQLRecordManager and PGVector, it is possible to guarantee that all processes will
be executed in a single transaction.

This is, moreover, the only solution we know of to guarantee the consistency of the
import of chunks into a vector database. It's possible only if the outer session is built
with the connection.

def main():
    db_url = "postgresql+psycopg://postgres:password_postgres@localhost:5432/"
    engine = create_engine(db_url, echo=True)
    embeddings = FakeEmbeddings()
    pgvector:VectorStore = PGVector(
        embeddings=embeddings,
        connection=engine,
    )

    record_manager = SQLRecordManager(
        namespace="namespace",
        engine=engine,
    )
    record_manager.create_schema()

    with engine.connect() as connection:
        session_maker = scoped_session(sessionmaker(bind=connection))
        # NOTE: Update session_factories
        record_manager.session_factory = session_maker
        pgvector.session_maker = session_maker
        with connection.begin():
            loader = CSVLoader(
                    "data/faq/faq.csv",
                    source_column="source",
                    autodetect_encoding=True,
                )
            result = index(
                source_id_key="source",
                docs_source=loader.load()[:1],
                cleanup="incremental",
                vector_store=pgvector,
                record_manager=record_manager,
            )
            print(result)

The same thing is possible asynchronously, but a bug in sql_record_manager.py
in _amake_session() must first be fixed (See PR ).

    async def _amake_session(self) -> AsyncGenerator[AsyncSession, None]:
        """Create a session and close it after use."""

        # FIXME: REMOVE if not isinstance(self.session_factory, async_sessionmaker):~~
        if not isinstance(self.engine, AsyncEngine):
            raise AssertionError("This method is not supported for sync engines.")

        async with self.session_factory() as session:
            yield session

Then, it is possible to do the same thing asynchronously:

async def main():
    db_url = "postgresql+psycopg://postgres:password_postgres@localhost:5432/"
    engine = create_async_engine(db_url, echo=True)
    embeddings = FakeEmbeddings()
    pgvector:VectorStore = PGVector(
        embeddings=embeddings,
        connection=engine,
    )
    record_manager = SQLRecordManager(
        namespace="namespace",
        engine=engine,
        async_mode=True,
    )
    await record_manager.acreate_schema()

    async with engine.connect() as connection:
        session_maker = async_scoped_session(
            async_sessionmaker(bind=connection),
            scopefunc=current_task)
        record_manager.session_factory = session_maker
        pgvector.session_maker = session_maker
        async with connection.begin():
            loader = CSVLoader(
                "data/faq/faq.csv",
                source_column="source",
                autodetect_encoding=True,
            )
            result = await aindex(
                source_id_key="source",
                docs_source=loader.load()[:1],
                cleanup="incremental",
                vector_store=pgvector,
                record_manager=record_manager,
            )
            print(result)


asyncio.run(main())

The promise of the constructor, with the create_extension parameter, is to guarantee that the extension is added before the APIs are used. Since this promise cannot be kept in an async scenario, there is an alternative:

Remove this parameter, since the promise cannot be kept. Otherwise, an async method is needed to install the extension before the APIs are used, and to check that this method has been invoked at the start of each API.
Use a lazy approach as suggested, which simply respects the constructor's promise.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/langchain-ai/langchain-postgres/pull/64 **Author:** [@pprados](https://github.com/pprados) **Created:** 6/10/2024 **Status:** ✅ Merged **Merged:** 6/10/2024 **Merged by:** [@eyurtsev](https://github.com/eyurtsev) **Base:** `main` ← **Head:** `pprados/async2` --- ### 📝 Commits (10+) - [`188c83d`](https://github.com/langchain-ai/langchain-postgres/commit/188c83dedd6ac605c6f9f2f4e1441a664563b0b9) Pre integration PPR - [`69c22a3`](https://github.com/langchain-ai/langchain-postgres/commit/69c22a3ce926f4b29d360d0b50724a0f89ed198f) Add async mode - [`0abd730`](https://github.com/langchain-ai/langchain-postgres/commit/0abd730671f35dd6ba148e7ee4ba1b2617780da2) Fix lint - [`e1ad8c8`](https://github.com/langchain-ai/langchain-postgres/commit/e1ad8c824283c83379d41b6d766a464524e4abd0) It's possible to share the session_maker with sync mode. - [`dc6187a`](https://github.com/langchain-ai/langchain-postgres/commit/dc6187a905cef3ce5218945ecae3595f7d23951e) Fix __apost_init__ - [`14d6da8`](https://github.com/langchain-ai/langchain-postgres/commit/14d6da818ca5fc50398057e93ee65e933627fac3) Add async API - [`e95df2f`](https://github.com/langchain-ai/langchain-postgres/commit/e95df2fe5e075008a7ef0ffe943d788b00eeb7e0) Rebase - [`9c2035c`](https://github.com/langchain-ai/langchain-postgres/commit/9c2035c3335e6c404902431c6505b18277179510) Fix create_vector_extension with async - [`7cc0bb1`](https://github.com/langchain-ai/langchain-postgres/commit/7cc0bb1f8d1dfdfa331548ca0f2866fb727cbb16) Fix create_vector_extension with async - [`4cbeeb2`](https://github.com/langchain-ai/langchain-postgres/commit/4cbeeb2772713a8301386752a0c366a9c6af91bc) Re-ajuste les UI vis à vis de l'async (partiellement) ### 📊 Changes **3 files changed** (+1328 additions, -86 deletions) <details> <summary>View changed files</summary> 📝 `langchain_postgres/vectorstores.py` (+830 -74) 📝 `pyproject.toml` (+4 -0) 📝 `tests/unit_tests/test_vectorstore.py` (+494 -12) </details> ### 📄 Description I made the mistake of doing a rebase while a review was in progress. This seems to block the process. There's a ‘1 change requested’ request that I can't validate. I propose another PR identical to the [previous one](https://github.com/langchain-ai/langchain-postgres/pull/32). This PR adds the **async** approach for pgvector. Some remarks: - We use assert to check invocations and not if. Thus, in production, it is possible to remove these checks with `python -O ...` - We propose a public `session_maker` attribute. This is very important for resilient scenarios. In a RAG architecture, it is necessary to import document chunks. To keep track of the links between chunks and documents, we can use the [index()](https://python.langchain.com/docs/modules/data_connection/indexing/) API. This API proposes to use an SQL-type record manager. In a classic use case, using `SQLRecordManager` and a vector database, it is impossible to guarantee the consistency of the import. Indeed, if a crash occurs during the import, there is an inconsistency between the SQL database and the vector database. **PGVector is the solution to this problem.** Indeed, it is possible to use a single database (and not a two-phase commit with 2 technologies, if they are both compatible). But, for this, it is necessary to be able to combine the transactions between the use of `SQLRecordManager` and `PGVector` as a vector database. This is only possible if it is possible to intervene on the `session_maker`. This is why we propose to make this attribute public. By unifying the `session_maker` of `SQLRecordManager` and `PGVector`, it is possible to guarantee that all processes will be executed in a single transaction. This is, moreover, the only solution we know of to guarantee the consistency of the import of chunks into a vector database. It's possible only if the outer session is built with the connection. ```python def main(): db_url = "postgresql+psycopg://postgres:password_postgres@localhost:5432/" engine = create_engine(db_url, echo=True) embeddings = FakeEmbeddings() pgvector:VectorStore = PGVector( embeddings=embeddings, connection=engine, ) record_manager = SQLRecordManager( namespace="namespace", engine=engine, ) record_manager.create_schema() with engine.connect() as connection: session_maker = scoped_session(sessionmaker(bind=connection)) # NOTE: Update session_factories record_manager.session_factory = session_maker pgvector.session_maker = session_maker with connection.begin(): loader = CSVLoader( "data/faq/faq.csv", source_column="source", autodetect_encoding=True, ) result = index( source_id_key="source", docs_source=loader.load()[:1], cleanup="incremental", vector_store=pgvector, record_manager=record_manager, ) print(result) ``` The same thing is possible asynchronously, but a bug in `sql_record_manager.py` in `_amake_session()` must first be fixed (See [PR](https://github.com/langchain-ai/langchain/pull/20735) ). ```python async def _amake_session(self) -> AsyncGenerator[AsyncSession, None]: """Create a session and close it after use.""" # FIXME: REMOVE if not isinstance(self.session_factory, async_sessionmaker):~~ if not isinstance(self.engine, AsyncEngine): raise AssertionError("This method is not supported for sync engines.") async with self.session_factory() as session: yield session ``` Then, it is possible to do the same thing asynchronously: ```python async def main(): db_url = "postgresql+psycopg://postgres:password_postgres@localhost:5432/" engine = create_async_engine(db_url, echo=True) embeddings = FakeEmbeddings() pgvector:VectorStore = PGVector( embeddings=embeddings, connection=engine, ) record_manager = SQLRecordManager( namespace="namespace", engine=engine, async_mode=True, ) await record_manager.acreate_schema() async with engine.connect() as connection: session_maker = async_scoped_session( async_sessionmaker(bind=connection), scopefunc=current_task) record_manager.session_factory = session_maker pgvector.session_maker = session_maker async with connection.begin(): loader = CSVLoader( "data/faq/faq.csv", source_column="source", autodetect_encoding=True, ) result = await aindex( source_id_key="source", docs_source=loader.load()[:1], cleanup="incremental", vector_store=pgvector, record_manager=record_manager, ) print(result) asyncio.run(main()) ``` The promise of the constructor, with the `create_extension` parameter, is to guarantee that the extension is added before the APIs are used. Since this promise cannot be kept in an `async` scenario, there is an alternative: - Remove this parameter, since the promise cannot be kept. Otherwise, an `async` method is needed to install the extension before the APIs are used, and to check that this method has been invoked at the start of each API. - Use a lazy approach as suggested, which simply respects the constructor's promise. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

yindo added the pull-request label 2026-02-16 05:16:44 -05:00

yindo closed this issue

2026-02-16 05:16:44 -05:00

Sign in to join this conversation.

1 Participants

Notifications

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: langchain-ai/langchain-postgres#142