Toggle Menu
Prev Up Next

# Source code for langchain_community.chat_message_histories.tidb

    import json
    import logging
    from datetime import datetime
    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 sqlalchemy import create_engine, text
    from sqlalchemy.exc import SQLAlchemyError
    from sqlalchemy.orm import sessionmaker

    logger = logging.getLogger(__name__)

    class TiDBChatMessageHistory(BaseChatMessageHistory):
        """
        Represents a chat message history stored in a TiDB database.
        """

        def __init__(
            self,
            session_id: str,
            connection_string: str,
            table_name: str = "langchain_message_store",
            earliest_time: Optional[datetime] = None,
        ):
            """
            Initializes a new instance of the TiDBChatMessageHistory class.

            Args:
                session_id (str): The ID of the chat session.
                connection_string (str): The connection string for the TiDB database.
                    format: mysql+pymysql://<host>:<PASSWORD>@<host>:4000/<db>?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true
                table_name (str, optional): the table name to store the chat messages.
                    Defaults to "langchain_message_store".
                earliest_time (Optional[datetime], optional): The earliest time to retrieve messages from.
                    Defaults to None.
            """  # noqa

            self.session_id = session_id
            self.table_name = table_name
            self.earliest_time = earliest_time
            self.cache: List = []

            # Set up SQLAlchemy engine and session
            self.engine = create_engine(connection_string)
            Session = sessionmaker(bind=self.engine)
            self.session = Session()

            self._create_table_if_not_exists()
            self._load_messages_to_cache()

        def _create_table_if_not_exists(self) -> None:
            """
            Creates a table if it does not already exist in the database.
            """

            create_table_query = text(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    session_id VARCHAR(255) NOT NULL,
                    message JSON NOT NULL,
                    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    INDEX session_idx (session_id)
                );"""
            )
            try:
                self.session.execute(create_table_query)
                self.session.commit()
            except SQLAlchemyError as e:
                logger.error(f"Error creating table: {e}")
                self.session.rollback()

        def _load_messages_to_cache(self) -> None:
            """
            Loads messages from the database into the cache.

            This method retrieves messages from the database table. The retrieved messages
            are then stored in the cache for faster access.

            Raises:
                SQLAlchemyError: If there is an error executing the database query.

            """
            time_condition = (
                f"AND create_time >= '{self.earliest_time}'" if self.earliest_time else ""
            )
            query = text(
                f"""
                SELECT message FROM {self.table_name}
                WHERE session_id = :session_id {time_condition}
                ORDER BY id;
            """
            )
            try:
                result = self.session.execute(query, {"session_id": self.session_id})
                for record in result.fetchall():
                    message_dict = json.loads(record[0])
                    self.cache.append(messages_from_dict([message_dict])[0])
            except SQLAlchemyError as e:
                logger.error(f"Error loading messages to cache: {e}")

        @property
        def messages(self) -> List[BaseMessage]:  # type: ignore[override]
            """returns all messages"""
            if len(self.cache) == 0:
                self.reload_cache()
            return self.cache

        def add_message(self, message: BaseMessage) -> None:
            """adds a message to the database and cache"""
            query = text(
                f"INSERT INTO {self.table_name} (session_id, message) VALUES (:session_id, :message);"  # noqa
            )
            try:
                self.session.execute(
                    query,
                    {
                        "session_id": self.session_id,
                        "message": json.dumps(message_to_dict(message)),
                    },
                )
                self.session.commit()
                self.cache.append(message)
            except SQLAlchemyError as e:
                logger.error(f"Error adding message: {e}")
                self.session.rollback()

        def clear(self) -> None:
            """clears all messages"""
            query = text(f"DELETE FROM {self.table_name} WHERE session_id = :session_id;")
            try:
                self.session.execute(query, {"session_id": self.session_id})
                self.session.commit()
                self.cache.clear()
            except SQLAlchemyError as e:
                logger.error(f"Error clearing messages: {e}")
                self.session.rollback()

        def reload_cache(self) -> None:
            """reloads messages from database to cache"""
            self.cache.clear()
            self._load_messages_to_cache()

        def __del__(self) -> None:
            """closes the session"""
            self.session.close()