Source code for langchain.memory.chat_message_histories.cassandra
import json
import logging
from typing import List
from langchain.schema import (
    BaseChatMessageHistory,
    BaseMessage,
    _message_to_dict,
    messages_from_dict,
)
logger = logging.getLogger(__name__)
DEFAULT_KEYSPACE_NAME = "chat_history"
DEFAULT_TABLE_NAME = "message_store"
DEFAULT_USERNAME = "cassandra"
DEFAULT_PASSWORD = "cassandra"
DEFAULT_PORT = 9042
[docs]class CassandraChatMessageHistory(BaseChatMessageHistory):
    """Chat message history that stores history in Cassandra.
    Args:
        contact_points: list of ips to connect to Cassandra cluster
        session_id: arbitrary key that is used to store the messages
            of a single chat session.
        port: port to connect to Cassandra cluster
        username: username to connect to Cassandra cluster
        password: password to connect to Cassandra cluster
        keyspace_name: name of the keyspace to use
        table_name: name of the table to use
    """
    def __init__(
        self,
        contact_points: List[str],
        session_id: str,
        port: int = DEFAULT_PORT,
        username: str = DEFAULT_USERNAME,
        password: str = DEFAULT_PASSWORD,
        keyspace_name: str = DEFAULT_KEYSPACE_NAME,
        table_name: str = DEFAULT_TABLE_NAME,
    ):
        self.contact_points = contact_points
        self.session_id = session_id
        self.port = port
        self.username = username
        self.password = password
        self.keyspace_name = keyspace_name
        self.table_name = table_name
        try:
            from cassandra import (
                AuthenticationFailed,
                OperationTimedOut,
                UnresolvableContactPoints,
            )
            from cassandra.cluster import Cluster, PlainTextAuthProvider
        except ImportError:
            raise ValueError(
                "Could not import cassandra-driver python package. "
                "Please install it with `pip install cassandra-driver`."
            )
        self.cluster: Cluster = Cluster(
            contact_points,
            port=port,
            auth_provider=PlainTextAuthProvider(
                username=self.username, password=self.password
            ),
        )
        try:
            self.session = self.cluster.connect()
        except (
            AuthenticationFailed,
            UnresolvableContactPoints,
            OperationTimedOut,
        ) as error:
            logger.error(
                "Unable to establish connection with \
                cassandra chat message history database"
            )
            raise error
        self._prepare_cassandra()
    def _prepare_cassandra(self) -> None:
        """Create the keyspace and table if they don't exist yet"""
        from cassandra import OperationTimedOut, Unavailable
        try:
            self.session.execute(
                f"""CREATE KEYSPACE IF NOT EXISTS 
                {self.keyspace_name} WITH REPLICATION = 
                {{ 'class' : 'SimpleStrategy', 'replication_factor' : 1 }};"""
            )
        except (OperationTimedOut, Unavailable) as error:
            logger.error(
                f"Unable to create cassandra \
                chat message history keyspace: {self.keyspace_name}."
            )
            raise error
        self.session.set_keyspace(self.keyspace_name)
        try:
            self.session.execute(
                f"""CREATE TABLE IF NOT EXISTS 
                {self.table_name} (id UUID, session_id varchar, 
                history text,  PRIMARY KEY ((session_id), id) );"""
            )
        except (OperationTimedOut, Unavailable) as error:
            logger.error(
                f"Unable to create cassandra \
                chat message history table: {self.table_name}"
            )
            raise error
    @property
    def messages(self) -> List[BaseMessage]:  # type: ignore
        """Retrieve the messages from Cassandra"""
        from cassandra import ReadFailure, ReadTimeout, Unavailable
        try:
            rows = self.session.execute(
                f"""SELECT * FROM {self.table_name}
                WHERE session_id = '{self.session_id}' ;"""
            )
        except (Unavailable, ReadTimeout, ReadFailure) as error:
            logger.error("Unable to Retreive chat history messages from cassadra")
            raise error
        if rows:
            items = [json.loads(row.history) for row in rows]
        else:
            items = []
        messages = messages_from_dict(items)
        return messages
[docs]    def add_message(self, message: BaseMessage) -> None:
        """Append the message to the record in Cassandra"""
        import uuid
        from cassandra import Unavailable, WriteFailure, WriteTimeout
        try:
            self.session.execute(
                """INSERT INTO message_store
                (id, session_id, history) VALUES (%s, %s, %s);""",
                (uuid.uuid4(), self.session_id, json.dumps(_message_to_dict(message))),
            )
        except (Unavailable, WriteTimeout, WriteFailure) as error:
            logger.error("Unable to write chat history messages to cassandra")
            raise error
[docs]    def clear(self) -> None:
        """Clear session memory from Cassandra"""
        from cassandra import OperationTimedOut, Unavailable
        try:
            self.session.execute(
                f"DELETE FROM {self.table_name} WHERE session_id = '{self.session_id}';"
            )
        except (Unavailable, OperationTimedOut) as error:
            logger.error("Unable to clear chat history messages from cassandra")
            raise error
    def __del__(self) -> None:
        if self.session:
            self.session.shutdown()
        if self.cluster:
            self.cluster.shutdown()