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.

[zach-hopkins]

This worked for my purposes in case it helps anybody out. Using a custom async class for RedisChatMessageHistory to use RunnableWithMessageHistory

import json
import logging
from typing import List, Optional, Sequence  # Keep Sequence for potential future use

import redis.asyncio as aioredis
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import (
    BaseMessage,
    message_to_dict,
    messages_from_dict,
)

logger = logging.getLogger(__name__)


class AsyncRedisChatMessageHistory(BaseChatMessageHistory):
    """
    Async Chat message history stored in a Redis list, using an externally
    managed async Redis client instance. Designed for async environments.
    Synchronous methods will raise errors if called from an active event loop.
    """

    def __init__(
        self,
        session_id: str,
        redis_client: aioredis.Redis,  # Expects an initialized async client
        key_prefix: str = "message_store:",
        ttl: Optional[int] = 60 * 60 * 24 * 7,  # Default TTL: 7 days
        max_length: Optional[int] = None,  # Default: No length limit
    ):
        """Initialize the history store."""
        # Basic check for required async methods
        if not all(
            hasattr(redis_client, method)
            for method in ["lpush", "lrange", "ltrim", "expire", "delete"]
        ):
            raise TypeError(
                "redis_client must be an aioredis.Redis instance with list commands."
            )

        self.redis_client = redis_client
        self.session_id = session_id
        self.key_prefix = key_prefix
        self.ttl = ttl
        self.max_length = max_length
        # No loop storage needed

    @property
    def key(self) -> str:
        """Constructs the Redis key for the session."""
        return self.key_prefix + self.session_id

    # --- Async Methods (Primary Interface) ---

    async def aget_messages(self) -> List[BaseMessage]:
        """Retrieve messages from Redis asynchronously."""
        try:
            # Assumes client uses decode_responses=True
            items_json = await self.redis_client.lrange(self.key, 0, -1)
            if not items_json:
                return []
        except Exception as e:
            logger.error(f"Redis error getting messages for key {self.key}: {e}")
            return []
        items = []
        for m_json in items_json:
            try:
                items.append(json.loads(m_json))
            except json.JSONDecodeError:
                logger.warning(f"Corrupt JSON in key {self.key}: {m_json[:100]}...")
        # LPUSH adds to left (newest = index 0). LRANGE gets [newest, ..., oldest]
        # messages_from_dict expects [oldest, ..., newest]. Reverse the list.
        try:
            messages = messages_from_dict(items[::-1])  # Reverse for correct order
        except Exception as e:
            logger.error(f"Error converting dicts to messages for key {self.key}: {e}")
            return []
        return messages

    async def aadd_message(self, message: BaseMessage) -> None:
        """Append the message to the record in Redis asynchronously."""
        try:
            # Assumes client uses decode_responses=True
            message_json = json.dumps(message_to_dict(message))
            # LPUSH adds the message to the beginning (left/head) of the list
            await self.redis_client.lpush(self.key, message_json)

            # Trim history if max_length is set
            if self.max_length is not None and self.max_length > 0:
                # Keep only the first (newest) max_length items
                await self.redis_client.ltrim(self.key, 0, self.max_length - 1)

            # Set/reset TTL if specified
            if self.ttl is not None and self.ttl > 0:
                await self.redis_client.expire(self.key, self.ttl)
            # Allow persisting if ttl is explicitly 0 or None (handled by default)

        except Exception as e:
            logger.error(f"Redis error adding message for key {self.key}: {e}")
            # Consider re-raising

    # Optional, but good practice for BaseChatMessageHistory
    async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
        """Add multiple messages async. Uses atomic transaction."""
        try:
            # Use Redis pipeline for atomic and efficient multi-message add
            async with self.redis_client.pipeline(transaction=True) as pipe:
                for message in messages:
                    message_json = json.dumps(message_to_dict(message))
                    # LPUSH adds to the head, so add them in reverse if order matters for reading
                    # Or add normally and handle order on read (as done in aget_messages)
                    pipe.lpush(self.key, message_json)

                # Optional Trim and TTL apply after messages are added
                if self.max_length is not None and self.max_length > 0:
                    pipe.ltrim(self.key, 0, self.max_length - 1)
                if self.ttl is not None and self.ttl > 0:
                    pipe.expire(self.key, self.ttl)

                await pipe.execute()
        except Exception as e:
            logger.error(
                f"Redis error adding multiple messages for key {self.key}: {e}"
            )
            # Consider re-raising

    async def aclear(self) -> None:
        """Clear session memory from Redis asynchronously."""
        try:
            await self.redis_client.delete(self.key)
        except Exception as e:
            logger.error(f"Redis error clearing key {self.key}: {e}")
            # Consider re-raising

    # --- Sync Methods (Implemented as non-blocking error stubs) ---

    @property
    def messages(self) -> List[BaseMessage]:  # type: ignore[override] # Keep ignore for mypy property override issue
        """Synchronous access is not supported. Use 'await history.aget_messages()'."""
        logger.error(
            "Synchronous 'messages' property was called on AsyncRedisChatMessageHistory!"
        )
        raise RuntimeError(
            "Synchronous access to 'messages' is not supported in this async "
            "implementation. Use 'await .aget_messages()' instead."
        )

    def add_message(self, message: BaseMessage) -> None:
        """Synchronous addition is not supported. Use 'await history.aadd_message()'."""
        logger.error(
            "Synchronous 'add_message' called on AsyncRedisChatMessageHistory!"
        )
        raise RuntimeError(
            "Synchronous 'add_message' is not supported in this async implementation. "
            "Use 'await .aadd_message()' or 'await .aadd_messages()' instead."
        )

    # Optional: Implement sync add_messages if Base class requires it concretely
    def add_messages(self, messages: Sequence[BaseMessage]) -> None:
        """Synchronous bulk addition is not supported. Use 'await history.aadd_messages()'."""
        logger.error(
            "Synchronous 'add_messages' called on AsyncRedisChatMessageHistory!"
        )
        raise RuntimeError(
            "Synchronous 'add_messages' is not supported in this async implementation. "
            "Use 'await .aadd_messages()' instead."
        )

    def clear(self) -> None:
        """Synchronous clear is not supported. Use 'await history.aclear()'."""
        logger.error("Synchronous 'clear' called on AsyncRedisChatMessageHistory!")
        raise RuntimeError(
            "Synchronous 'clear' is not supported in this async implementation. "
            "Use 'await .aclear()' instead."
        )