{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zep\n",
    "## Retriever Example for [Zep](https://docs.getzep.com/) - A long-term memory store for LLM applications.\n",
    "\n",
    "### More on Zep:\n",
    "\n",
    "Zep stores, summarizes, embeds, indexes, and enriches conversational AI chat histories, and exposes them via simple, low-latency APIs.\n",
    "\n",
    "Key Features:\n",
    "\n",
    "- **Fast!** Zep’s async extractors operate independently of the your chat loop, ensuring a snappy user experience.\n",
    "- **Long-term memory persistence**, with access to historical messages irrespective of your summarization strategy.\n",
    "- **Auto-summarization** of memory messages based on a configurable message window. A series of summaries are stored, providing flexibility for future summarization strategies.\n",
    "- **Hybrid search** over memories and metadata, with messages automatically embedded on creation.\n",
    "- **Entity Extractor** that automatically extracts named entities from messages and stores them in the message metadata.\n",
    "- **Auto-token counting** of memories and summaries, allowing finer-grained control over prompt assembly.\n",
    "- Python and JavaScript SDKs.\n",
    "\n",
    "Zep project: [https://github.com/getzep/zep](https://github.com/getzep/zep)\n",
    "Docs: [https://docs.getzep.com/](https://docs.getzep.com/)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retriever Example\n",
    "\n",
    "This notebook demonstrates how to search historical chat message histories using the [Zep Long-term Memory Store](https://getzep.github.io/).\n",
    "\n",
    "We'll demonstrate:\n",
    "\n",
    "1. Adding conversation history to the Zep memory store.\n",
    "2. Vector search over the conversation history.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-25T15:03:27.863217Z",
     "start_time": "2023-05-25T15:03:25.690273Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from langchain.memory.chat_message_histories import ZepChatMessageHistory\n",
    "from langchain.schema import HumanMessage, AIMessage\n",
    "from uuid import uuid4\n",
    "import getpass\n",
    "\n",
    "# Set this to your Zep server URL\n",
    "ZEP_API_URL = \"http://localhost:8000\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize the Zep Chat Message History Class and add a chat message history to the memory store\n",
    "\n",
    "**NOTE:** Unlike other Retrievers, the content returned by the Zep Retriever is session/user specific. A `session_id` is required when instantiating the Retriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      " ········\n"
     ]
    }
   ],
   "source": [
    "# Provide your Zep API key. Note that this is optional. See https://docs.getzep.com/deployment/auth\n",
    "\n",
    "zep_api_key = getpass.getpass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-25T15:03:29.118416Z",
     "start_time": "2023-05-25T15:03:29.022464Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "session_id = str(uuid4())  # This is a unique identifier for the user/session\n",
    "\n",
    "# Set up Zep Chat History. We'll use this to add chat histories to the memory store\n",
    "zep_chat_history = ZepChatMessageHistory(\n",
    "    session_id=session_id, url=ZEP_API_URL, api_key=zep_api_key\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-25T15:03:30.271181Z",
     "start_time": "2023-05-25T15:03:30.180442Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Preload some messages into the memory. The default message window is 12 messages. We want to push beyond this to demonstrate auto-summarization.\n",
    "test_history = [\n",
    "    {\"role\": \"human\", \"content\": \"Who was Octavia Butler?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"Octavia Estelle Butler (June 22, 1947 – February 24, 2006) was an American\"\n",
    "            \" science fiction author.\"\n",
    "        ),\n",
    "    },\n",
    "    {\"role\": \"human\", \"content\": \"Which books of hers were made into movies?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"The most well-known adaptation of Octavia Butler's work is the FX series\"\n",
    "            \" Kindred, based on her novel of the same name.\"\n",
    "        ),\n",
    "    },\n",
    "    {\"role\": \"human\", \"content\": \"Who were her contemporaries?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"Octavia Butler's contemporaries included Ursula K. Le Guin, Samuel R.\"\n",
    "            \" Delany, and Joanna Russ.\"\n",
    "        ),\n",
    "    },\n",
    "    {\"role\": \"human\", \"content\": \"What awards did she win?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"Octavia Butler won the Hugo Award, the Nebula Award, and the MacArthur\"\n",
    "            \" Fellowship.\"\n",
    "        ),\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"human\",\n",
    "        \"content\": \"Which other women sci-fi writers might I want to read?\",\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": \"You might want to read Ursula K. Le Guin or Joanna Russ.\",\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"human\",\n",
    "        \"content\": (\n",
    "            \"Write a short synopsis of Butler's book, Parable of the Sower. What is it\"\n",
    "            \" about?\"\n",
    "        ),\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"Parable of the Sower is a science fiction novel by Octavia Butler,\"\n",
    "            \" published in 1993. It follows the story of Lauren Olamina, a young woman\"\n",
    "            \" living in a dystopian future where society has collapsed due to\"\n",
    "            \" environmental disasters, poverty, and violence.\"\n",
    "        ),\n",
    "    },\n",
    "]\n",
    "\n",
    "for msg in test_history:\n",
    "    zep_chat_history.add_message(\n",
    "        HumanMessage(content=msg[\"content\"])\n",
    "        if msg[\"role\"] == \"human\"\n",
    "        else AIMessage(content=msg[\"content\"])\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use the Zep Retriever to vector search over the Zep memory\n",
    "\n",
    "Zep provides native vector search over historical conversation memory. Embedding happens automatically.\n",
    "\n",
    "NOTE: Embedding of messages occurs asynchronously, so the first query may not return results. Subsequent queries will return results as the embeddings are generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-25T15:03:32.979155Z",
     "start_time": "2023-05-25T15:03:32.590310Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='Parable of the Sower is a science fiction novel by Octavia Butler, published in 1993. It follows the story of Lauren Olamina, a young woman living in a dystopian future where society has collapsed due to environmental disasters, poverty, and violence.', metadata={'score': 0.8897116216176073, 'uuid': 'db60ff57-f259-4ec4-8a81-178ed4c6e54f', 'created_at': '2023-06-26T23:40:22.816214Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'GPE', 'Matches': [{'End': 20, 'Start': 15, 'Text': 'Sower'}], 'Name': 'Sower'}, {'Label': 'PERSON', 'Matches': [{'End': 65, 'Start': 51, 'Text': 'Octavia Butler'}], 'Name': 'Octavia Butler'}, {'Label': 'DATE', 'Matches': [{'End': 84, 'Start': 80, 'Text': '1993'}], 'Name': '1993'}, {'Label': 'PERSON', 'Matches': [{'End': 124, 'Start': 110, 'Text': 'Lauren Olamina'}], 'Name': 'Lauren Olamina'}]}}, 'token_count': 56}),\n",
       " Document(page_content=\"Write a short synopsis of Butler's book, Parable of the Sower. What is it about?\", metadata={'score': 0.8856661080361157, 'uuid': 'f1a5981a-8f6d-4168-a548-6e9c32f35fa1', 'created_at': '2023-06-26T23:40:22.809621Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'ORG', 'Matches': [{'End': 32, 'Start': 26, 'Text': 'Butler'}], 'Name': 'Butler'}, {'Label': 'WORK_OF_ART', 'Matches': [{'End': 61, 'Start': 41, 'Text': 'Parable of the Sower'}], 'Name': 'Parable of the Sower'}]}}, 'token_count': 23}),\n",
       " Document(page_content='Who was Octavia Butler?', metadata={'score': 0.7757595298492976, 'uuid': '361d0043-1009-4e13-a7f0-8aea8b1ee869', 'created_at': '2023-06-26T23:40:22.709886Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'PERSON', 'Matches': [{'End': 22, 'Start': 8, 'Text': 'Octavia Butler'}], 'Name': 'Octavia Butler'}], 'intent': 'The subject wants to know about the identity or background of an individual named Octavia Butler.'}}, 'token_count': 8}),\n",
       " Document(page_content=\"Octavia Butler's contemporaries included Ursula K. Le Guin, Samuel R. Delany, and Joanna Russ.\", metadata={'score': 0.7601242516059306, 'uuid': '56c45e8a-0f65-45f0-bc46-d9e65164b563', 'created_at': '2023-06-26T23:40:22.778836Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'PERSON', 'Matches': [{'End': 16, 'Start': 0, 'Text': \"Octavia Butler's\"}], 'Name': \"Octavia Butler's\"}, {'Label': 'ORG', 'Matches': [{'End': 58, 'Start': 41, 'Text': 'Ursula K. Le Guin'}], 'Name': 'Ursula K. Le Guin'}, {'Label': 'PERSON', 'Matches': [{'End': 76, 'Start': 60, 'Text': 'Samuel R. Delany'}], 'Name': 'Samuel R. Delany'}, {'Label': 'PERSON', 'Matches': [{'End': 93, 'Start': 82, 'Text': 'Joanna Russ'}], 'Name': 'Joanna Russ'}], 'intent': \"The subject is providing information about Octavia Butler's contemporaries.\"}}, 'token_count': 27}),\n",
       " Document(page_content='You might want to read Ursula K. Le Guin or Joanna Russ.', metadata={'score': 0.7594731095320668, 'uuid': '6951f2fd-dfa4-4e05-9380-f322ef8f72f8', 'created_at': '2023-06-26T23:40:22.80464Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'ORG', 'Matches': [{'End': 40, 'Start': 23, 'Text': 'Ursula K. Le Guin'}], 'Name': 'Ursula K. Le Guin'}, {'Label': 'PERSON', 'Matches': [{'End': 55, 'Start': 44, 'Text': 'Joanna Russ'}], 'Name': 'Joanna Russ'}]}}, 'token_count': 18})]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.retrievers import ZepRetriever\n",
    "\n",
    "zep_retriever = ZepRetriever(\n",
    "    session_id=session_id,  # Ensure that you provide the session_id when instantiating the Retriever\n",
    "    url=ZEP_API_URL,\n",
    "    top_k=5,\n",
    "    api_key=zep_api_key,\n",
    ")\n",
    "\n",
    "await zep_retriever.aget_relevant_documents(\"Who wrote Parable of the Sower?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the Zep sync API to retrieve results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-25T15:03:34.713354Z",
     "start_time": "2023-05-25T15:03:34.577974Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='Parable of the Sower is a science fiction novel by Octavia Butler, published in 1993. It follows the story of Lauren Olamina, a young woman living in a dystopian future where society has collapsed due to environmental disasters, poverty, and violence.', metadata={'score': 0.889661105796371, 'uuid': 'db60ff57-f259-4ec4-8a81-178ed4c6e54f', 'created_at': '2023-06-26T23:40:22.816214Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'GPE', 'Matches': [{'End': 20, 'Start': 15, 'Text': 'Sower'}], 'Name': 'Sower'}, {'Label': 'PERSON', 'Matches': [{'End': 65, 'Start': 51, 'Text': 'Octavia Butler'}], 'Name': 'Octavia Butler'}, {'Label': 'DATE', 'Matches': [{'End': 84, 'Start': 80, 'Text': '1993'}], 'Name': '1993'}, {'Label': 'PERSON', 'Matches': [{'End': 124, 'Start': 110, 'Text': 'Lauren Olamina'}], 'Name': 'Lauren Olamina'}]}}, 'token_count': 56}),\n",
       " Document(page_content=\"Write a short synopsis of Butler's book, Parable of the Sower. What is it about?\", metadata={'score': 0.885754241595424, 'uuid': 'f1a5981a-8f6d-4168-a548-6e9c32f35fa1', 'created_at': '2023-06-26T23:40:22.809621Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'ORG', 'Matches': [{'End': 32, 'Start': 26, 'Text': 'Butler'}], 'Name': 'Butler'}, {'Label': 'WORK_OF_ART', 'Matches': [{'End': 61, 'Start': 41, 'Text': 'Parable of the Sower'}], 'Name': 'Parable of the Sower'}]}}, 'token_count': 23}),\n",
       " Document(page_content='Who was Octavia Butler?', metadata={'score': 0.7758688965570713, 'uuid': '361d0043-1009-4e13-a7f0-8aea8b1ee869', 'created_at': '2023-06-26T23:40:22.709886Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'PERSON', 'Matches': [{'End': 22, 'Start': 8, 'Text': 'Octavia Butler'}], 'Name': 'Octavia Butler'}], 'intent': 'The subject wants to know about the identity or background of an individual named Octavia Butler.'}}, 'token_count': 8}),\n",
       " Document(page_content=\"Octavia Butler's contemporaries included Ursula K. Le Guin, Samuel R. Delany, and Joanna Russ.\", metadata={'score': 0.7602672137411663, 'uuid': '56c45e8a-0f65-45f0-bc46-d9e65164b563', 'created_at': '2023-06-26T23:40:22.778836Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'PERSON', 'Matches': [{'End': 16, 'Start': 0, 'Text': \"Octavia Butler's\"}], 'Name': \"Octavia Butler's\"}, {'Label': 'ORG', 'Matches': [{'End': 58, 'Start': 41, 'Text': 'Ursula K. Le Guin'}], 'Name': 'Ursula K. Le Guin'}, {'Label': 'PERSON', 'Matches': [{'End': 76, 'Start': 60, 'Text': 'Samuel R. Delany'}], 'Name': 'Samuel R. Delany'}, {'Label': 'PERSON', 'Matches': [{'End': 93, 'Start': 82, 'Text': 'Joanna Russ'}], 'Name': 'Joanna Russ'}], 'intent': \"The subject is providing information about Octavia Butler's contemporaries.\"}}, 'token_count': 27}),\n",
       " Document(page_content='You might want to read Ursula K. Le Guin or Joanna Russ.', metadata={'score': 0.7596040989115522, 'uuid': '6951f2fd-dfa4-4e05-9380-f322ef8f72f8', 'created_at': '2023-06-26T23:40:22.80464Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'ORG', 'Matches': [{'End': 40, 'Start': 23, 'Text': 'Ursula K. Le Guin'}], 'Name': 'Ursula K. Le Guin'}, {'Label': 'PERSON', 'Matches': [{'End': 55, 'Start': 44, 'Text': 'Joanna Russ'}], 'Name': 'Joanna Russ'}]}}, 'token_count': 18})]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zep_retriever.get_relevant_documents(\"Who wrote Parable of the Sower?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-18T20:09:21.298710Z",
     "start_time": "2023-05-18T20:09:21.297169Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}