community[patch]: Make sql record manager fully compatible with async

[pprados]

The _amake_session() method does not allow modifying the self.session_factory with
anything other than async_sessionmaker. This prohibits advanced uses of index().

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
(problem with the network, ...)
there is an inconsistency between the SQL database and the vector database.

With the PR we are proposing for langchain-postgres,
it is now possible 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.

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

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Ignored Deployment

Name	Status	Preview	Comments	Updated (UTC)
langchain	⬜️ Ignored (Inspect)	Visit Preview		May 6, 2024 11:35am

[eyurtsev]

This code is mutating in place attributes on the classes.

        record_manager.session_factory = session_maker
        pgvector.session_maker = session_maker

While attributes are not marked explicitly as private, it's safest to still treat these as such.

---

Are you able to pass required information via the initializer? If so, what information do you need?

    def __init__(
        self,
        namespace: str,
        *,
        connection: Optional[Union[??]], # <-- Is there something that can be provided here that will accommodate your use use case,
        engine: Optional[Union[Engine, AsyncEngine]] = None,
        db_url: Union[None, str, URL] = None,
        engine_kwargs: Optional[Dict[str, Any]] = None,
        async_mode: bool = False,
    ) -> None:

[pprados]

The SQLRecordManager can not be initialized with connection. And, it's not enough.

The session_maker = scoped_session(sessionmaker(bind=connection)) and async_scoped_session(async_sessionmaker(bind=connection),scopefunc=current_task) must be used by all SQL tools.
May be, it's possible to add a parameter with current session_maker for PGVector and SQLRecordManager.

It's a complex process. The session_maker MUST be initialized with the CURRENT connection (one connection by thread or async job. It's not possible to share the connection with different threads. It's not reentrant). Only with this construction, it's possible to share all the sub-transaction with the outer transaction, and guarantee full execution of the import, whatever happens.

The LCEL approch propose to use a singleton to describe the chain. So, you must use an Engine, and not a Connection. The connection is open on when it's need, one for each job.

The attributes record_manager.session_factory and pgvector.session_maker may be updated for that.

[pprados]

Hello from Paris 😄

To fully grasp the difficulty of this, one must understand that scoped_session() is a class, which has, as an attribute, a local() allowing variables associated with the current thread. Therefore, within the same thread, multiple instances of scoped_session() can exist, each having its own local variable to keep the session active. Hence, if scoped_session() is used in both the PGVector() and SQLRecordManager() classes, they will not share the same session, leading to inconsistency in case of a crash during import.

One could have a global variable global_scoped_session with an instance of scoped_session() for langchain, and ensure that all SQLAlchemy APIs use this global variable for session creation. Here, similarly, it's the caller, responsible for creating the first session, who can decide how to create it.

However, this doesn't work for our use case because, during the initialization of the global_scoped_session variable, a sessionmaker(bind=connection) needs to be created, and it's not possible to use a connection for a global variable. Possibly, a sessionmaker(bind=engine), but it's strange to have an engine imposed by langchain globally.

One could imagine proposing an alternative to scoped_session that uses a global variable for local(): global_scoped_session. Thus, if both classes default to using global_scoped_session(), if the invocation is made directly, a session is created. There's no longer a need to modify session_(maker|factory).

If, before the call, global_scoped_session() is invoked, PGVector and SQLRecordManager use this session. Therefore, it's the caller who can decide how the original session will be created.

So, to conclude, the only viable solution, in my opinion, is the one I propose. By offering the ability to modify the session_factory and session_maker, my PR can be validated 😄.

[pprados]

Hello @eyurtsev

Can you reconsider this very important small PR? Without it, it's not possible to have a resilient RAG in production.

[eyurtsev]

Shouldn't be included

[eyurtsev]

Why doesn't the original code work?