langchain[minor], community[minor], core[minor]: Async Cache support and AsyncRedisCache

[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]

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

[cbornet]

An alternative would be to add these methods to the BaseCache class with a default implementation that calls the non-async method with run_in_executor like what is done in BaseLLM and Embeddings.
Otherwise, there is no reason to prefix the methods here with the a prefix.

[eyurtsev]

@dzmitry-kankalovich We generally prefer keeping async version of methods in the same class as the sync versions, and have them provide a default async implementation that uses run_in_executor like @cbornet suggested.

Would you mind merging into existing abstractions? Ideally a single PR that just modifies the core interface first, and then separately we can do a PR for any implementations

[dzmitry-kankalovich]

Done

[dzmitry-kankalovich]

As for the smaller PRs - yes, I can move slower and split up into several PRs, if you are happy with current direction.

[eyurtsev]

Generally better to split PRs by package to minimize potential dependency conflicts since the packages may have different release schedules

[eyurtsev]

@dzmitry-kankalovich thanks for the contribution! LangChain uses the convention of having sync and async implementations defined in the same abstraction

[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.

[cbornet]

Maybe for this you could have a RedisCache that uses a redis.Redis client and an AsyncRedisCache that uses a redis.asyncio.Redis ? (I know this is the opposite move as for BaseCache 😄 )

[dzmitry-kankalovich]

I actually did it this way in the very first variant, but I didn't like it was literally 99% of the copy from the original RedisCache. The remaining 1% were async/await keywords and client check in constructor. Feels like a potential future maintenance problem?

[dzmitry-kankalovich]

Also, I am not sure what I suppose to do for AsyncRedisCache sync methods lookup/update/clear. Like I can raise error on those, or do the opposite trick for async in sync execution, but doesn't it feel like too much of duplication of effort?

[cbornet]

You could put this common behavior in a parent abstract class ?

[dzmitry-kankalovich]

Alright, I can do it. But before that, can you tell me whats your rationale here? Just so I know for future :)

[cbornet]

RedisBaseCache is abstract and has methods _key, __ensure_generation_type __get_generations __configure_pipeline_for_update
RedisCache extends RedisBaseCache and implements methods lookup, update, clear using redis.Redis
AsyncRedisCache extends RedisBaseCache and implements methods alookup, aupdate, aclear using redis.asyncio.Redis
AsyncRedisCache raises NotImplementedError() for lookup, update, clear

[dzmitry-kankalovich]

@cbornet I mean I was asking whats the motivation for such split.

I've pushed the implementation which is basically what you said, however I went a bit further with sync-in-async and actually made it so you can run it in certain context (weird case though, I log warning for that). I can rework to NotImplementedError if you think its cleaner.

[eyurtsev]

or do the opposite trick for async in sync execution

At the moment, we generally try to avoid that. Only one event loop is allowed in a given thread, so there's some additional logic that's required to determine if there's an event loop running already and if yes, kick off another thread etc

[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.

[cbornet]

I'm not sure it's a good thing to output these logs. It will flood the logs.

[dzmitry-kankalovich]

agree.

[cbornet]

I think this is too complex and brittle. I would just throw NotImplementedError and nothing else.

[dzmitry-kankalovich]

agree.

[dzmitry-kankalovich]

@cbornet adjusted code per your review.

[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]

@dzmitry-kankalovich I added a docker compose file under /docker we might move it at some point, but thinking about making it easier for developers to spin up services that integration tests depend on

[dzmitry-kankalovich]

yeah that would be super useful!

[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).

[baskaryan]

note: shouldn't import this in langchain, imports here are only for backwards compatibility

[dzmitry-kankalovich]

I see. Noted! Maybe makes sense to drop a comment in that file next time you are working with it - future contributors will know.