{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "91c6a7ef",
   "metadata": {},
   "source": [
    "# Dynamodb Chat Message History\n",
    "\n",
    "This notebook goes over how to use Dynamodb to store chat message history."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f608be0",
   "metadata": {},
   "source": [
    "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 boto3."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "030d784f",
   "metadata": {},
   "source": [
    "Next, create the DynamoDB Table where we will be storing messages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "id": "d15e3302",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory.chat_message_histories import DynamoDBChatMessageHistory\n",
    "\n",
    "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": 3,
   "id": "64fc465e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='hi!', additional_kwargs={}, example=False),\n",
       " AIMessage(content='whats up?', additional_kwargs={}, example=False)]"
      ]
     },
     "execution_count": 3,
     "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": null,
   "id": "225713c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory.chat_message_histories import DynamoDBChatMessageHistory\n",
    "\n",
    "history = DynamoDBChatMessageHistory(\n",
    "    table_name=\"SessionTable\",\n",
    "    session_id=\"0\",\n",
    "    endpoint_url=\"http://localhost.localstack.cloud:4566\",\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3b33c988",
   "metadata": {},
   "source": [
    "## Agent with DynamoDB Memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f92d9499",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import Tool\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.agents import initialize_agent\n",
    "from langchain.agents import AgentType\n",
    "from langchain.utilities import PythonREPL\n",
    "from getpass import getpass\n",
    "\n",
    "message_history = DynamoDBChatMessageHistory(table_name=\"SessionTable\", session_id=\"1\")\n",
    "memory = ConversationBufferMemory(\n",
    "    memory_key=\"chat_history\", chat_memory=message_history, return_messages=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1167eeba",
   "metadata": {},
   "outputs": [],
   "source": [
    "python_repl = PythonREPL()\n",
    "\n",
    "# You can create the tool to pass to an agent\n",
    "tools = [\n",
    "    Tool(\n",
    "        name=\"python_repl\",\n",
    "        description=\"A Python shell. Use this to execute python commands. Input should be a valid python command. If you want to see the output of a value, you should print it out with `print(...)`.\",\n",
    "        func=python_repl.run,\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fce085c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(temperature=0)\n",
    "agent_chain = initialize_agent(\n",
    "    tools,\n",
    "    llm,\n",
    "    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,\n",
    "    verbose=True,\n",
    "    memory=memory,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "952a3103",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m{\n",
      "    \"action\": \"Final Answer\",\n",
      "    \"action_input\": \"Hello! How can I assist you today?\"\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Hello! How can I assist you today?'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_chain.run(input=\"Hello!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "54c4aaf4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m{\n",
      "    \"action\": \"python_repl\",\n",
      "    \"action_input\": \"import requests\\nfrom bs4 import BeautifulSoup\\n\\nurl = 'https://en.wikipedia.org/wiki/Twitter'\\nresponse = requests.get(url)\\nsoup = BeautifulSoup(response.content, 'html.parser')\\nowner = soup.find('th', text='Owner').find_next_sibling('td').text.strip()\\nprint(owner)\"\n",
      "}\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mX Corp. (2023–present)Twitter, Inc. (2006–2023)\n",
      "\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m{\n",
      "    \"action\": \"Final Answer\",\n",
      "    \"action_input\": \"X Corp. (2023–present)Twitter, Inc. (2006–2023)\"\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'X Corp. (2023–present)Twitter, Inc. (2006–2023)'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_chain.run(input=\"Who owns Twitter?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f9013118",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m{\n",
      "    \"action\": \"Final Answer\",\n",
      "    \"action_input\": \"Hello Bob! How can I assist you today?\"\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Hello Bob! How can I assist you today?'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_chain.run(input=\"My name is Bob.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "405e5315",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m{\n",
      "    \"action\": \"Final Answer\",\n",
      "    \"action_input\": \"Your name is Bob.\"\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Your name is Bob.'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_chain.run(input=\"Who am I?\")"
   ]
  }
 ],
 "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}