Source code for langchain.memory.chat_message_histories.rocksetdb
from datetime import datetime
from time import sleep
from typing import Any, Callable, List, Union
from uuid import uuid4
from langchain.schema import BaseChatMessageHistory
from langchain.schema.messages import BaseMessage, _message_to_dict, messages_from_dict
[docs]class RocksetChatMessageHistory(BaseChatMessageHistory):
    """Uses Rockset to store chat messages.
    To use, ensure that the `rockset` python package installed.
    Example:
        .. code-block:: python
            from langchain.memory.chat_message_histories import (
                RocksetChatMessageHistory
            )
            from rockset import RocksetClient
            history = RocksetChatMessageHistory(
                session_id="MySession",
                client=RocksetClient(),
                collection="langchain_demo",
                sync=True
            )
            history.add_user_message("hi!")
            history.add_ai_message("whats up?")
            print(history.messages)
    """
    # You should set these values based on your VI.
    # These values are configured for the typical
    # free VI. Read more about VIs here:
    # https://rockset.com/docs/instances
    SLEEP_INTERVAL_MS: int = 5
    ADD_TIMEOUT_MS: int = 5000
    CREATE_TIMEOUT_MS: int = 20000
    def _wait_until(self, method: Callable, timeout: int, **method_params: Any) -> None:
        """Sleeps until meth() evaluates to true. Passes kwargs into
        meth.
        """
        start = datetime.now()
        while not method(**method_params):
            curr = datetime.now()
            if (curr - start).total_seconds() * 1000 > timeout: