langchain, community, core: Async Cache support introduced

[dzmitry-kankalovich]

Description

We're using FastAPI in our current project with async functions from LangChain. Recently we needed to add Redis-based cache. However we quickly learned that regardless of the LangChain way of usage - sync or async - the cache lookup/update/clear will happen always in synchronous manner. Naturally, this means blocking the event loop, especially with slower than Redis cache implementations.

So I took a similar approach to agenerate and ainvoke for LLMs, and created an AsyncBaseCache mixin interface with corresponding alookup/aupdate/aclear async variants:

class AsyncBaseCache(ABC):
    """Base interface for async cache."""

    @abstractmethod
    async def alookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]:
        """Look up based on prompt and llm_string."""

    @abstractmethod
    async def aupdate(
        self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE
    ) -> None:
        """Update cache based on prompt and llm_string."""

    @abstractmethod
    async def aclear(self, **kwargs: Any) -> None:
        """Clear cache that can take additional keyword arguments."""

The reference implementation is done for InMemoryCache and RedisCache.

The InMemoryCache does not really need async variant, since there is no real IO, however having that this is the most convenient way to unit test changes.

The change spans throughout core, community and langchain. Linter, formatters, tests are passing.

Issue

Fixes inability to use IO-bound Redis cache w/o blocking event loop. Opens opportunity for other cache implementations to have async variants.

Dependencies

None

Twitter

@Mind_Clash

[vercel[bot]]

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		Feb 8, 2024 3:01am

[dzmitry-kankalovich]

@cbornet @eyurtsev thanks, I will look into @cbornet idea and see if I can make it work

[dzmitry-kankalovich]

@cbornet @eyurtsev I've pushed reworked interfaces - thanks for suggestion, it now feels much simpler.

I do however run into unit test problem with SQLite-based cache, which I cannot figure out yet. Somehow the base aupdate with default exec context does not work well for SQLite cache, although I've tested these changes with my production code and see no issues.

Anyway, I am going to figure out whats the problem and get back to you.

UPDATE: It's a problem specifically with in-memory SQLite used in tests.

UPDATE 2: problem fixed.

[dzmitry-kankalovich]

@cbornet @eyurtsev I think at this point I've addressed all the current review feedback - please check it out, if you've got time.

The only remained suggestion is to have smaller PRs for this - I can split this one up, if you're happy with overall approach.

[dzmitry-kankalovich]

@cbornet adjusted code per your review.

[vercel[bot]]

Deployment failed with the following error:

The provided GitHub repository does not contain the requested branch or commit reference. Please ensure the repository is not empty.

[cbornet]

LGTM, great work !

[dzmitry-kankalovich]

@eyurtsev I think I am done as far as the changes go. If you have some time - let me know how we can proceed forward. Thanks!

[eyurtsev]

@eyurtsev I think I am done as far as the changes go. If you have some time - let me know how we can proceed forward. Thanks!

Apologies for taking so long to review! I'll try to follow up quickly on any changes

[eyurtsev]

@baskaryan FYI -- change to the base schema, so we'll need to bump minimal dependency in langchain. Do you need me to break these into two separate PRs or are we okay merging, releasing core, bumping constraints in langchain, and then releasing langchain?

[baskaryan]

@baskaryan FYI -- change to the base schema, so we'll need to bump minimal dependency in langchain. Do you need me to break these into two separate PRs or are we okay merging, releasing core, bumping constraints in langchain, and then releasing langchain?

can land in one 👍

[dzmitry-kankalovich]

@eyurtsev @baskaryan: It's all ready on my end, I don't plan to add anything else, it's good to go.

[dzmitry-kankalovich]

@hwchase17 @baskaryan hey guys, was this accidentally closed?

[baskaryan]

not intentional, we can't reopen at the moment but will do so as soon as we can :/ See #16796

[eyurtsev]

Taking a look now to see if I can refactor into a single class

[eyurtsev]

• Refactored to a single class
• Added a docker-compose file to make it easier to launch dependencies for tests - This needs to be documented in developer's guide and more services added
  • Dependencies will use non standard ports
• Refactored some of the cache tests to use context managers to make sure that the caches are always flushed between unit tests

---

I have to resolve a few mypy issues looks like mypy has a hard time discriminating between the async and the sync clients

[eyurtsev]

Decided to revert for now. Will merge as is!

The implementation with one class is more compact (https://github.com/langchain-ai/langchain/pull/15817/files#diff-60f65acc09d072aa20f3bad2cce0ecaee7e4ded197a56cd54d6671d76a684de0R401), but a bit harder to understand, and dealing with mypy is seems pretty annoying in that case.

I fixed a few typing issues in the original PR (missing type annotations).