{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f22eab3f84cbeb37",
   "metadata": {},
   "source": [
    "# SQL (SQLAlchemy)\n",
    "\n",
    ">[Structured Query Language (SQL)](https://en.wikipedia.org/wiki/SQL) is a domain-specific language used in programming and designed for managing data held in a relational database management system (RDBMS), or for stream processing in a relational data stream management system (RDSMS). It is particularly useful in handling structured data, i.e., data incorporating relations among entities and variables.\n",
    "\n",
    ">[SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy) is an open-source `SQL` toolkit and object-relational mapper (ORM) for the Python programming language released under the MIT License.\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."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8f2830ee9ca1e01",
   "metadata": {},
   "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.\n",
    "3. Install `SQLAlchemy` python package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab016290-3823-4e1b-9610-ae9a1b71cb07",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install SQLAlchemy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4576e914a866fb40",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-28T10:04:38.077748Z",
     "start_time": "2023-08-28T10:04:36.105894Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from langchain.memory.chat_message_histories import SQLChatMessageHistory\n",
    "\n",
    "chat_message_history = SQLChatMessageHistory(\n",
    "    session_id=\"test_session\", connection_string=\"sqlite:///sqlite.db\"\n",
    ")\n",
    "\n",
    "chat_message_history.add_user_message(\"Hello\")\n",
    "chat_message_history.add_ai_message(\"Hi\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b476688cbb32ba90",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-28T10:04:38.929396Z",
     "start_time": "2023-08-28T10:04:38.915727Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e5337719d5614fd",
   "metadata": {},
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fdfde84c07d071bb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-28T10:04:41.510498Z",
     "start_time": "2023-08-28T10:04:41.494912Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "from typing import Any\n",
    "\n",
    "from langchain.memory.chat_message_histories.sql import BaseMessageConverter\n",
    "from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage\n",
    "from sqlalchemy import Column, DateTime, Integer, Text\n",
    "from sqlalchemy.orm import declarative_base\n",
    "\n",
    "Base = declarative_base()\n",
    "\n",
    "\n",
    "class CustomMessage(Base):\n",
    "    __tablename__ = \"custom_message_store\"\n",
    "\n",
    "    id = Column(Integer, primary_key=True)\n",
    "    session_id = Column(Text)\n",
    "    type = Column(Text)\n",
    "    content = Column(Text)\n",
    "    created_at = Column(DateTime)\n",
    "    author_email = Column(Text)\n",
    "\n",
    "\n",
    "class CustomMessageConverter(BaseMessageConverter):\n",
    "    def __init__(self, author_email: str):\n",
    "        self.author_email = author_email\n",
    "\n",
    "    def from_sql_model(self, sql_message: Any) -> BaseMessage:\n",
    "        if sql_message.type == \"human\":\n",
    "            return HumanMessage(\n",
    "                content=sql_message.content,\n",
    "            )\n",
    "        elif sql_message.type == \"ai\":\n",
    "            return AIMessage(\n",
    "                content=sql_message.content,\n",
    "            )\n",
    "        elif sql_message.type == \"system\":\n",
    "            return SystemMessage(\n",
    "                content=sql_message.content,\n",
    "            )\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown message type: {sql_message.type}\")\n",
    "\n",
    "    def to_sql_model(self, message: BaseMessage, session_id: str) -> Any:\n",
    "        now = datetime.now()\n",
    "        return CustomMessage(\n",
    "            session_id=session_id,\n",
    "            type=message.type,\n",
    "            content=message.content,\n",
    "            created_at=now,\n",
    "            author_email=self.author_email,\n",
    "        )\n",
    "\n",
    "    def get_sql_model_class(self) -> Any:\n",
    "        return CustomMessage\n",
    "\n",
    "\n",
    "chat_message_history = SQLChatMessageHistory(\n",
    "    session_id=\"test_session\",\n",
    "    connection_string=\"sqlite:///sqlite.db\",\n",
    "    custom_message_converter=CustomMessageConverter(author_email=\"test@example.com\"),\n",
    ")\n",
    "\n",
    "chat_message_history.add_user_message(\"Hello\")\n",
    "chat_message_history.add_ai_message(\"Hi\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4a6a54d8a9e2856f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-28T10:04:43.497990Z",
     "start_time": "2023-08-28T10:04:43.492517Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "622aded629a1adeb",
   "metadata": {},
   "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": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
