Should we add an asynchronous RedisChatMessageHistory implemetion?

[zhangkai803]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Take advantage of asynchronous capabilities when using RedisChatMessageHistory.

Motivation

Since Runnable provides some asynchronous functions, we'd better support asynchrony in the underlying storage.
As mentioned in #9100, but no solution worked.

Proposal (If applicable)

Here is an implemention with aioredis based on #1895 :

import aioredis
import asyncio
import json
from typing import List, Optional

from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, message_to_dict, messages_from_dict


class AsyncRedisChatMessageHistory(BaseChatMessageHistory):

    def __init__(
        self,
        session_id: str,
        url: str = "redis://localhost:6379/0",
        key_prefix: str = "message_store:",
        ttl: Optional[int] = None,
        loop = None,
    ):

        try:
            import aioredis
        except ImportError:
            raise ImportError(
                "Could not import aioredis python package. "
                "Please install it with `pip install aioredis`."
            )

        try:
            self.redis_client = aioredis.Redis.from_url(url)
        except aioredis.ConnectionError as error:
            logger.error(error)

        self.session_id = session_id
        self.key_prefix = key_prefix
        self.ttl = ttl
        self.loop = loop if loop is not None else asyncio.get_event_loop()

    @property
    def key(self) -> str:
        return self.key_prefix + self.session_id

    @property
    def messages(self):
        return asyncio.run_coroutine_threadsafe(self.aget_messages(), self.loop).result()

    async def aget_messages(self) -> List[BaseMessage]:  # type: ignore
        _items = await self.redis_client.lrange(self.key, 0, -1)
        items = [json.loads(m.decode("utf-8")) for m in _items[::-1]]
        messages = messages_from_dict(items)
        return messages

    def add_message(self, message: BaseMessage) -> None:
        asyncio.run_coroutine_threadsafe(self.aadd_message(message), self.loop)

    async def aadd_message(self, message: BaseMessage) -> None:
        await self.redis_client.lpush(self.key, json.dumps(message_to_dict(message)))
        if self.ttl:
            await self.redis_client.expire(self.key, self.ttl)

    def clear(self):
        asyncio.run_coroutine_threadsafe(self.aclear(), self.loop)

    async def aclear(self) -> None:
        await self.redis_client.delete(self.key)

Here is a demo:

import aioredis
import asyncio
import json
from typing import List, Optional

from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, message_to_dict, messages_from_dict
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI


class AsyncRedisChatMessageHistory(BaseChatMessageHistory):

    def __init__(
        self,
        session_id: str,
        url: str = "redis://localhost:6379/0",
        key_prefix: str = "message_store:",
        ttl: Optional[int] = None,
    ):
        self.redis_client = aioredis.Redis.from_url(url)
        self.session_id = session_id
        self.key_prefix = key_prefix
        self.ttl = ttl
        self.loop = asyncio.get_event_loop()

    @property
    def key(self) -> str:
        return self.key_prefix + self.session_id

    @property
    def messages(self):
        return asyncio.run_coroutine_threadsafe(self.aget_messages(), self.loop).result()

    async def aget_messages(self) -> List[BaseMessage]:  # type: ignore
        _items = await self.redis_client.lrange(self.key, 0, -1)
        items = [json.loads(m.decode("utf-8")) for m in _items[::-1]]
        messages = messages_from_dict(items)
        return messages

    def add_message(self, message: BaseMessage) -> None:
        asyncio.run_coroutine_threadsafe(self.aadd_message(message), self.loop)

    async def aadd_message(self, message: BaseMessage) -> None:
        await self.redis_client.lpush(self.key, json.dumps(message_to_dict(message)))
        if self.ttl:
            await self.redis_client.expire(self.key, self.ttl)

    def clear(self):
        asyncio.run_coroutine_threadsafe(self.aclear(), self.loop)

    async def aclear(self) -> None:
        await self.redis_client.delete(self.key)


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    return AsyncRedisChatMessageHistory(session_id, ttl=600)


model = ChatOpenAI()
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You're an assistant who's good at {ability}. Respond in 20 words or fewer.",
        ),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ]
)
runnable = prompt | model

with_message_history = RunnableWithMessageHistory(
    runnable,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history",
)


async def main():
    print(await with_message_history.ainvoke(
        {"ability": "math", "input": "What does cosine mean?"},
        config={"configurable": {"session_id": "abc123"}},
    ))

    await asyncio.gather(
        # Remembers
        with_message_history.ainvoke(
            {"ability": "math", "input": f"What?"},
            config={"configurable": {"session_id": "abc123"}},
        ),
        # New session_id --> does not remember.
        with_message_history.ainvoke(
            {"ability": "math", "input": f"What?"},
            config={"configurable": {"session_id": "def234"}},
        )
    )

    # Wait for the last add_message to be done
    await asyncio.sleep(3)


if __name__ == "__main__":
    asyncio.get_event_loop().run_until_complete(main())

[diego898]

@zhangkai803 - This is great! I was just looking for this - what about this?

https://redis-py.readthedocs.io/en/stable/examples/asyncio_examples.html

I just had to code up message truncation with redis, and now need it to be async as well

[zhangkai803]

If you want to use redis-py instead of aioredis:

• Replace import aioredis with import redis.asyncio as redis at line1.
• Replace self.redis_client = aioredis.Redis.from_url(url) with self.redis_client = redis.from_url(url) at line22.

[diego898]

Yeah that’s what I did - aioredis became redis-py I believe.

Just wanted to support this request and suggest it gets merged!

[diego898]

Note that RunnableWithMessage history's _exit_history function does not asyncronous write, even if the BaseChatMessageHistory has support for asynchronous writing.