{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "91c6a7ef",
   "metadata": {},
   "source": [
    "# AWS DynamoDB\n",
    "\n",
    ">[Amazon AWS DynamoDB](https://awscli.amazonaws.com/v2/documentation/api/latest/reference/dynamodb/index.html) is a fully managed `NoSQL` database service that provides fast and predictable performance with seamless scalability.\n",
    "\n",
    "This notebook goes over how to use `DynamoDB` to store chat message history with `DynamoDBChatMessageHistory` class."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bcbd170",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First make sure you have correctly configured the [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html). Then make sure you have installed the `langchain-community` package, so we need to install that. We also need to install the `boto3` package.\n",
    "\n",
    "```bash\n",
    "pip install -U langchain-community boto3\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdec429d",
   "metadata": {},
   "source": [
    "It's also helpful (but not needed) to set up [LangSmith](https://smith.langchain.com/) for best-in-class observability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47d3f725",
   "metadata": {},
   "outputs": [],
   "source": [
    "# os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "# os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c68f434-0ccf-49a4-8313-e9f6064086a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.chat_message_histories import (\n",
    "    DynamoDBChatMessageHistory,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "030d784f",
   "metadata": {},
   "source": [
    "## Create Table\n",
    "\n",
    "Now, create the `DynamoDB` Table where we will be storing messages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "93ce1811",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "import boto3\n",
    "\n",
    "# Get the service resource.\n",
    "dynamodb = boto3.resource(\"dynamodb\")\n",
    "\n",
    "# Create the DynamoDB table.\n",
    "table = dynamodb.create_table(\n",
    "    TableName=\"SessionTable\",\n",
    "    KeySchema=[{\"AttributeName\": \"SessionId\", \"KeyType\": \"HASH\"}],\n",
    "    AttributeDefinitions=[{\"AttributeName\": \"SessionId\", \"AttributeType\": \"S\"}],\n",
    "    BillingMode=\"PAY_PER_REQUEST\",\n",
    ")\n",
    "\n",
    "# Wait until the table exists.\n",
    "table.meta.client.get_waiter(\"table_exists\").wait(TableName=\"SessionTable\")\n",
    "\n",
    "# Print out some data about the table.\n",
    "print(table.item_count)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a9b310b",
   "metadata": {},
   "source": [
    "## DynamoDBChatMessageHistory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d15e3302",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = DynamoDBChatMessageHistory(table_name=\"SessionTable\", session_id=\"0\")\n",
    "\n",
    "history.add_user_message(\"hi!\")\n",
    "\n",
    "history.add_ai_message(\"whats up?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "64fc465e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='hi!'), AIMessage(content='whats up?')]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "955f1b15",
   "metadata": {},
   "source": [
    "## DynamoDBChatMessageHistory with Custom Endpoint URL\n",
    "\n",
    "Sometimes it is useful to specify the URL to the AWS endpoint to connect to. For instance, when you are running locally against [Localstack](https://localstack.cloud/). For those cases you can specify the URL via the `endpoint_url` parameter in the constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "225713c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = DynamoDBChatMessageHistory(\n",
    "    table_name=\"SessionTable\",\n",
    "    session_id=\"0\",\n",
    "    endpoint_url=\"http://localhost.localstack.cloud:4566\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97f8578a",
   "metadata": {},
   "source": [
    "## DynamoDBChatMessageHistory With Composite Keys\n",
    "The default key for DynamoDBChatMessageHistory is ```{\"SessionId\": self.session_id}```, but you can modify this to match your table design.\n",
    "\n",
    "### Primary Key Name\n",
    "You may modify the primary key by passing in a primary_key_name value in the constructor, resulting in the following:\n",
    "```{self.primary_key_name: self.session_id}```\n",
    "\n",
    "### Composite Keys\n",
    "When using an existing DynamoDB table, you may need to modify the key structure from the default of to something including a Sort Key. To do this you may use the ```key``` parameter.\n",
    "\n",
    "Passing a value for key will override the primary_key parameter, and the resulting key structure will be the passed value.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "088c037c",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "composite_table = dynamodb.create_table(\n",
    "    TableName=\"CompositeTable\",\n",
    "    KeySchema=[\n",
    "        {\"AttributeName\": \"PK\", \"KeyType\": \"HASH\"},\n",
    "        {\"AttributeName\": \"SK\", \"KeyType\": \"RANGE\"},\n",
    "    ],\n",
    "    AttributeDefinitions=[\n",
    "        {\"AttributeName\": \"PK\", \"AttributeType\": \"S\"},\n",
    "        {\"AttributeName\": \"SK\", \"AttributeType\": \"S\"},\n",
    "    ],\n",
    "    BillingMode=\"PAY_PER_REQUEST\",\n",
    ")\n",
    "\n",
    "# Wait until the table exists.\n",
    "composite_table.meta.client.get_waiter(\"table_exists\").wait(TableName=\"CompositeTable\")\n",
    "\n",
    "# Print out some data about the table.\n",
    "print(composite_table.item_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f462660f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='hello, composite dynamodb table!')]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_key = {\n",
    "    \"PK\": \"session_id::0\",\n",
    "    \"SK\": \"langchain_history\",\n",
    "}\n",
    "\n",
    "composite_key_history = DynamoDBChatMessageHistory(\n",
    "    table_name=\"CompositeTable\",\n",
    "    session_id=\"0\",\n",
    "    endpoint_url=\"http://localhost.localstack.cloud:4566\",\n",
    "    key=my_key,\n",
    ")\n",
    "\n",
    "composite_key_history.add_user_message(\"hello, composite dynamodb table!\")\n",
    "\n",
    "composite_key_history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131e42d2",
   "metadata": {},
   "source": [
    "## Chaining\n",
    "\n",
    "We can easily combine this message history class with [LCEL Runnables](/docs/expression_language/how_to/message_history)\n",
    "\n",
    "To do this we will want to use OpenAI, so we need to install that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbd0e91e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8cecb08",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant.\"),\n",
    "        MessagesPlaceholder(variable_name=\"history\"),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chain = prompt | ChatOpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88ec540f",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain_with_history = RunnableWithMessageHistory(\n",
    "    chain,\n",
    "    lambda session_id: DynamoDBChatMessageHistory(\n",
    "        table_name=\"SessionTable\", session_id=session_id\n",
    "    ),\n",
    "    input_messages_key=\"question\",\n",
    "    history_messages_key=\"history\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e9969fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is where we configure the session id\n",
    "config = {\"configurable\": {\"session_id\": \"<SESSION_ID>\"}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eb73f547",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Hello Bob! How can I assist you today?')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain_with_history.invoke({\"question\": \"Hi! I'm bob\"}, config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7daa3508",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Your name is Bob! Is there anything specific you would like assistance with, Bob?')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain_with_history.invoke({\"question\": \"Whats my name\"}, config=config)"
   ]
  }
 ],
 "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
}