codekingpro/portable-devtools

1import json2import logging3from typing import List, Optional4 5from langchain_core.chat_history import BaseChatMessageHistory6from langchain_core.messages import (7    BaseMessage,8    message_to_dict,9    messages_from_dict,10)11 12from langchain_community.utilities.redis import get_client13 14logger = logging.getLogger(__name__)15 16 17class RedisChatMessageHistory(BaseChatMessageHistory):18    """Chat message history stored in a Redis database.19 20    Setup:21        Install ``redis`` python package.22 23        .. code-block:: bash24 25            pip install redis26 27    Instantiate:28        .. code-block:: python29 30        from langchain_community.chat_message_histories import RedisChatMessageHistory31 32        history = RedisChatMessageHistory(33            session_id = "your-session-id",34            url="redis://your-host:your-port:your-database",  # redis://localhost:6379/035        )36 37    Add and retrieve messages:38        .. code-block:: python39 40            # Add single message41            history.add_message(message)42 43            # Add batch messages44            history.add_messages([message1, message2, message3, ...])45 46            # Add human message47            history.add_user_message(human_message)48 49            # Add ai message50            history.add_ai_message(ai_message)51 52            # Retrieve messages53            messages = history.messages54    """  # noqa: E50155 56    def __init__(57        self,58        session_id: str,59        url: str = "redis://localhost:6379/0",60        key_prefix: str = "message_store:",61        ttl: Optional[int] = None,62    ):63        """Initialize with a RedisChatMessageHistory instance.64 65        Args:66            session_id: str67                The ID for single chat session. Used to form keys with `key_prefix`.68            url: Optional[str]69                String parameter configuration for connecting to the redis.70            key_prefix: Optional[str]71                The prefix of the key, combined with `session id` to form the key.72            ttl: Optional[int]73                Set the expiration time of `key`, the unit is seconds.74        """75        try:76            import redis77        except ImportError:78            raise ImportError(79                "Could not import redis python package. "80                "Please install it with `pip install redis`."81            )82 83        try:84            self.redis_client = get_client(redis_url=url)85        except redis.exceptions.ConnectionError as error:86            logger.error(error)87 88        self.session_id = session_id89        self.key_prefix = key_prefix90        self.ttl = ttl91 92    @property93    def key(self) -> str:94        """Construct the record key to use"""95        return self.key_prefix + self.session_id96 97    @property98    def messages(self) -> List[BaseMessage]:99        """Retrieve the messages from Redis"""100        _items = self.redis_client.lrange(self.key, 0, -1)101        items = [json.loads(m.decode("utf-8")) for m in _items[::-1]]102        messages = messages_from_dict(items)103        return messages104 105    @messages.setter106    def messages(self, messages: List[BaseMessage]) -> None:107        raise NotImplementedError(108            "Direct assignment to 'messages' is not allowed."109            " Use the 'add_messages' instead."110        )111 112    def add_message(self, message: BaseMessage) -> None:113        """Append the message to the record in Redis"""114        self.redis_client.lpush(self.key, json.dumps(message_to_dict(message)))115        if self.ttl:116            self.redis_client.expire(self.key, self.ttl)117 118    def clear(self) -> None:119        """Clear session memory from Redis"""120        self.redis_client.delete(self.key)121