{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# SQL Chat Message History\n",
    "\n",
    "This notebook goes over a **SQLChatMessageHistory** class that allows to store chat history in any database supported by SQLAlchemy.\n",
    "\n",
    "Please note that to use it with databases other than SQLite, you will need to install the corresponding database driver."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f22eab3f84cbeb37"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Basic Usage\n",
    "\n",
    "To use the storage you need to provide only 2 things:\n",
    "\n",
    "1. Session Id - a unique identifier of the session, like user name, email, chat id etc.\n",
    "2. Connection string - a string that specifies the database connection. It will be passed to SQLAlchemy create_engine function."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f8f2830ee9ca1e01"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "from langchain.memory.chat_message_histories import SQLChatMessageHistory\n",
    "\n",
    "chat_message_history = SQLChatMessageHistory(\n",
    "\tsession_id='test_session',\n",
    "\tconnection_string='sqlite:///sqlite.db'\n",
    ")\n",
    "\n",
    "chat_message_history.add_user_message('Hello')\n",
    "chat_message_history.add_ai_message('Hi')"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-08-28T10:04:38.077748Z",
     "start_time": "2023-08-28T10:04:36.105894Z"
    }
   },
   "id": "4576e914a866fb40"
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "data": {
      "text/plain": "[HumanMessage(content='Hello', additional_kwargs={}, example=False),\n AIMessage(content='Hi', additional_kwargs={}, example=False)]"
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_message_history.messages"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-08-28T10:04:38.929396Z",
     "start_time": "2023-08-28T10:04:38.915727Z"
    }
   },
   "id": "b476688cbb32ba90"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Custom Storage Format\n",
    "\n",
    "By default, only the session id and message dictionary are stored in the table.\n",
    "\n",
    "However, sometimes you might want to store some additional information, like message date, author, language etc.\n",
    "\n",
    "To do that, you can create a custom message converter, by implementing **BaseMessageConverter** interface."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2e5337719d5614fd"
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "from langchain.schema import BaseMessage, HumanMessage, AIMessage, SystemMessage\n",
    "from typing import Any\n",
    "from sqlalchemy import Column, Integer, Text, DateTime\n",
    "from sqlalchemy.orm import declarative_base\n",
    "from langchain.memory.chat_message_histories.sql import BaseMessageConverter\n",
    "\n",
    "\n",
    "Base = declarative_base()\n",
    "\n",
    "\n",
    "class CustomMessage(Base):\n",
    "\t__tablename__ = 'custom_message_store'\n",
    "\n",
    "\tid = Column(Integer, primary_key=True)\n",
    "\tsession_id = Column(Text)\n",
    "\ttype = Column(Text)\n",
    "\tcontent = Column(Text)\n",
    "\tcreated_at = Column(DateTime)\n",
    "\tauthor_email = Column(Text)\n",
    "\n",
    "\n",
    "class CustomMessageConverter(BaseMessageConverter):\n",
    "\tdef __init__(self, author_email: str):\n",
    "\t\tself.author_email = author_email\n",
    "\t\n",
    "\tdef from_sql_model(self, sql_message: Any) -> BaseMessage:\n",
    "\t\tif sql_message.type == 'human':\n",
    "\t\t\treturn HumanMessage(\n",
    "\t\t\t\tcontent=sql_message.content,\n",
    "\t\t\t)\n",
    "\t\telif sql_message.type == 'ai':\n",
    "\t\t\treturn AIMessage(\n",
    "\t\t\t\tcontent=sql_message.content,\n",
    "\t\t\t)\n",
    "\t\telif sql_message.type == 'system':\n",
    "\t\t\treturn SystemMessage(\n",
    "\t\t\t\tcontent=sql_message.content,\n",
    "\t\t\t)\n",
    "\t\telse:\n",
    "\t\t\traise ValueError(f'Unknown message type: {sql_message.type}')\n",
    "\t\n",
    "\tdef to_sql_model(self, message: BaseMessage, session_id: str) -> Any:\n",
    "\t\tnow = datetime.now()\n",
    "\t\treturn CustomMessage(\n",
    "\t\t\tsession_id=session_id,\n",
    "\t\t\ttype=message.type,\n",
    "\t\t\tcontent=message.content,\n",
    "\t\t\tcreated_at=now,\n",
    "\t\t\tauthor_email=self.author_email\n",
    "\t\t)\n",
    "\t\n",
    "\tdef get_sql_model_class(self) -> Any:\n",
    "\t\treturn CustomMessage\n",
    "\n",
    "\n",
    "chat_message_history = SQLChatMessageHistory(\n",
    "\tsession_id='test_session',\n",
    "\tconnection_string='sqlite:///sqlite.db',\n",
    "\tcustom_message_converter=CustomMessageConverter(\n",
    "\t\tauthor_email='test@example.com'\n",
    "    )\n",
    ")\n",
    "\n",
    "chat_message_history.add_user_message('Hello')\n",
    "chat_message_history.add_ai_message('Hi')"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-08-28T10:04:41.510498Z",
     "start_time": "2023-08-28T10:04:41.494912Z"
    }
   },
   "id": "fdfde84c07d071bb"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "data": {
      "text/plain": "[HumanMessage(content='Hello', additional_kwargs={}, example=False),\n AIMessage(content='Hi', additional_kwargs={}, example=False)]"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_message_history.messages"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-08-28T10:04:43.497990Z",
     "start_time": "2023-08-28T10:04:43.492517Z"
    }
   },
   "id": "4a6a54d8a9e2856f"
  },
  {
   "cell_type": "markdown",
   "source": [
    "You also might want to change the name of session_id column. In this case you'll need to specify `session_id_field_name` parameter."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "622aded629a1adeb"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
