{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eg0Hwptz9g5q"
   },
   "source": [
    "# Entity Memory with SQLite storage\n",
    "\n",
    "In this walkthrough we'll create a simple conversation chain which uses ConversationEntityMemory backed by a SqliteEntityStore."
   ],
   "id": "d464a12a"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "2wUMSUoF8ffn"
   },
   "outputs": [],
   "source": [
    "from langchain.chains import ConversationChain\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.memory import ConversationEntityMemory\n",
    "from langchain.memory.entity import SQLiteEntityStore\n",
    "from langchain.memory.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE"
   ],
   "id": "db59b901"
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "8TpJZti99gxV"
   },
   "outputs": [],
   "source": [
    "entity_store = SQLiteEntityStore()\n",
    "llm = OpenAI(temperature=0)\n",
    "memory = ConversationEntityMemory(llm=llm, entity_store=entity_store)\n",
    "conversation = ConversationChain(\n",
    "    llm=llm,\n",
    "    prompt=ENTITY_MEMORY_CONVERSATION_TEMPLATE,\n",
    "    memory=memory,\n",
    "    verbose=True,\n",
    ")"
   ],
   "id": "ca6dee29"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HEAHG1L79ca1"
   },
   "source": [
    "Notice the usage of `EntitySqliteStore` as parameter to `entity_store` on the `memory` property."
   ],
   "id": "f9b4c3a0"
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 437
    },
    "id": "BzXphJWf_TAZ",
    "outputId": "de7fc966-e0fd-4daf-a9bd-4743455ea774"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant to a human, powered by a large language model trained by OpenAI.\n",
      "\n",
      "You are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
      "\n",
      "You are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
      "\n",
      "Overall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\n",
      "\n",
      "Context:\n",
      "{'Deven': 'Deven is working on a hackathon project with Sam.', 'Sam': 'Sam is working on a hackathon project with Deven.'}\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Last line:\n",
      "Human: Deven & Sam are working on a hackathon project\n",
      "You:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' That sounds like a great project! What kind of project are they working on?'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.run(\"Deven & Sam are working on a hackathon project\")"
   ],
   "id": "297e78a6"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "YsFE3hBjC6gl",
    "outputId": "56ab5ca9-e343-41b5-e69d-47541718a9b4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Deven is working on a hackathon project with Sam.'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.memory.entity_store.get(\"Deven\")"
   ],
   "id": "7e71f1dc"
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Sam is working on a hackathon project with Deven.'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.memory.entity_store.get(\"Sam\")"
   ],
   "id": "316f2e8d"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [],
   "id": "b85f8427"
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "venv"
  },
  "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
}