codekingpro/portable-devtools

1"""Chat message history stores a history of the message interactions in a chat."""2 3from __future__ import annotations4 5from abc import ABC, abstractmethod6from typing import TYPE_CHECKING7 8from pydantic import BaseModel, Field9 10from langchain_core.messages import (11    AIMessage,12    BaseMessage,13    HumanMessage,14    get_buffer_string,15)16from langchain_core.runnables.config import run_in_executor17 18if TYPE_CHECKING:19    from collections.abc import Sequence20 21 22class BaseChatMessageHistory(ABC):23    """Abstract base class for storing chat message history.24 25    Implementations guidelines:26 27    Implementations are expected to over-ride all or some of the following methods:28 29    * `add_messages`: sync variant for bulk addition of messages30    * `aadd_messages`: async variant for bulk addition of messages31    * `messages`: sync variant for getting messages32    * `aget_messages`: async variant for getting messages33    * `clear`: sync variant for clearing messages34    * `aclear`: async variant for clearing messages35 36    `add_messages` contains a default implementation that calls `add_message`37    for each message in the sequence. This is provided for backwards compatibility38    with existing implementations which only had `add_message`.39 40    Async variants all have default implementations that call the sync variants.41    Implementers can choose to override the async implementations to provide42    truly async implementations.43 44    Usage guidelines:45 46    When used for updating history, users should favor usage of `add_messages`47    over `add_message` or other variants like `add_user_message` and `add_ai_message`48    to avoid unnecessary round-trips to the underlying persistence layer.49 50    Example:51        ```python52        import json53        import os54        from langchain_core.messages import messages_from_dict, message_to_dict55 56 57        class FileChatMessageHistory(BaseChatMessageHistory):58            storage_path: str59            session_id: str60 61            @property62            def messages(self) -> list[BaseMessage]:63                try:64                    with open(65                        os.path.join(self.storage_path, self.session_id),66                        "r",67                        encoding="utf-8",68                    ) as f:69                        messages_data = json.load(f)70                    return messages_from_dict(messages_data)71                except FileNotFoundError:72                    return []73 74            def add_messages(self, messages: Sequence[BaseMessage]) -> None:75                all_messages = list(self.messages)  # Existing messages76                all_messages.extend(messages)  # Add new messages77 78                serialized = [message_to_dict(message) for message in all_messages]79                file_path = os.path.join(self.storage_path, self.session_id)80                os.makedirs(os.path.dirname(file_path), exist_ok=True)81                with open(file_path, "w", encoding="utf-8") as f:82                    json.dump(serialized, f)83 84            def clear(self) -> None:85                file_path = os.path.join(self.storage_path, self.session_id)86                os.makedirs(os.path.dirname(file_path), exist_ok=True)87                with open(file_path, "w", encoding="utf-8") as f:88                    json.dump([], f)89        ```90    """91 92    messages: list[BaseMessage]93    """A property or attribute that returns a list of messages.94 95    In general, getting the messages may involve IO to the underlying persistence96    layer, so this operation is expected to incur some latency.97    """98 99    async def aget_messages(self) -> list[BaseMessage]:100        """Async version of getting messages.101 102        Can over-ride this method to provide an efficient async implementation.103 104        In general, fetching messages may involve IO to the underlying persistence105        layer.106 107        Returns:108            The messages.109        """110        return await run_in_executor(None, lambda: self.messages)111 112    def add_user_message(self, message: HumanMessage | str) -> None:113        """Convenience method for adding a human message string to the store.114 115        !!! note116 117            This is a convenience method. Code should favor the bulk `add_messages`118            interface instead to save on round-trips to the persistence layer.119 120        This method may be deprecated in a future release.121 122        Args:123            message: The `HumanMessage` to add to the store.124        """125        if isinstance(message, HumanMessage):126            self.add_message(message)127        else:128            self.add_message(HumanMessage(content=message))129 130    def add_ai_message(self, message: AIMessage | str) -> None:131        """Convenience method for adding an `AIMessage` string to the store.132 133        !!! note134 135            This is a convenience method. Code should favor the bulk `add_messages`136            interface instead to save on round-trips to the persistence layer.137 138        This method may be deprecated in a future release.139 140        Args:141            message: The `AIMessage` to add.142        """143        if isinstance(message, AIMessage):144            self.add_message(message)145        else:146            self.add_message(AIMessage(content=message))147 148    def add_message(self, message: BaseMessage) -> None:149        """Add a Message object to the store.150 151        Args:152            message: A `BaseMessage` object to store.153 154        Raises:155            NotImplementedError: If the sub-class has not implemented an efficient156                `add_messages` method.157        """158        if type(self).add_messages != BaseChatMessageHistory.add_messages:159            # This means that the sub-class has implemented an efficient add_messages160            # method, so we should use it.161            self.add_messages([message])162        else:163            msg = (164                "add_message is not implemented for this class. "165                "Please implement add_message or add_messages."166            )167            raise NotImplementedError(msg)168 169    def add_messages(self, messages: Sequence[BaseMessage]) -> None:170        """Add a list of messages.171 172        Implementations should over-ride this method to handle bulk addition of messages173        in an efficient manner to avoid unnecessary round-trips to the underlying store.174 175        Args:176            messages: A sequence of `BaseMessage` objects to store.177        """178        for message in messages:179            self.add_message(message)180 181    async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:182        """Async add a list of messages.183 184        Args:185            messages: A sequence of `BaseMessage` objects to store.186        """187        await run_in_executor(None, self.add_messages, messages)188 189    @abstractmethod190    def clear(self) -> None:191        """Remove all messages from the store."""192 193    async def aclear(self) -> None:194        """Async remove all messages from the store."""195        await run_in_executor(None, self.clear)196 197    def __str__(self) -> str:198        """Return a string representation of the chat history."""199        return get_buffer_string(self.messages)200 201 202class InMemoryChatMessageHistory(BaseChatMessageHistory, BaseModel):203    """In memory implementation of chat message history.204 205    Stores messages in a memory list.206    """207 208    messages: list[BaseMessage] = Field(default_factory=list)209    """A list of messages stored in memory."""210 211    async def aget_messages(self) -> list[BaseMessage]:212        """Async version of getting messages.213 214        Can over-ride this method to provide an efficient async implementation.215 216        In general, fetching messages may involve IO to the underlying persistence217        layer.218 219        Returns:220            List of messages.221        """222        return self.messages223 224    def add_message(self, message: BaseMessage) -> None:225        """Add a self-created message to the store.226 227        Args:228            message: The message to add.229        """230        self.messages.append(message)231 232    async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:233        """Async add messages to the store.234 235        Args:236            messages: The messages to add.237        """238        self.add_messages(messages)239 240    def clear(self) -> None:241        """Clear all messages from the store."""242        self.messages = []243 244    async def aclear(self) -> None:245        """Async clear all messages from the store."""246        self.clear()247