langchain/docs/docs/integrations/retrievers/zep_memorystore.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zep\n",
    "## Retriever Example for [Zep](https://docs.getzep.com/)\n",
    "\n",
    "### Fast, Scalable Building Blocks for LLM Apps\n",
    "Zep is an open source platform for productionizing LLM apps. Go from a prototype\n",
    "built in LangChain or LlamaIndex, or a custom app, to production in minutes without\n",
    "rewriting code.\n",
    "\n",
    "Key Features:\n",
    "\n",
    "- **Fast!** Zep operates independently of the your chat loop, ensuring a snappy user experience.\n",
    "- **Chat History Memory, Archival, and Enrichment**, populate your prompts with relevant chat history, sumamries, named entities, intent data, and more.\n",
    "- **Vector Search over Chat History and Documents** Automatic embedding of documents, chat histories, and summaries. Use Zep's similarity or native MMR Re-ranked search to find the most relevant.\n",
    "- **Manage Users and their Chat Sessions** Users and their Chat Sessions are first-class citizens in Zep, allowing you to manage user interactions with your bots or agents easily.\n",
    "- **Records Retention and Privacy Compliance** Comply with corporate and regulatory mandates for records retention while ensuring compliance with privacy regulations such as CCPA and GDPR. Fulfill *Right To Be Forgotten* requests with a single API call\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"
   ]
  },
  {
   "attachments": {},
   "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",
    "    1. With a similarity search over chat messages\n",
    "    2. Using maximal marginal relevance re-ranking of a chat message search\n",
    "    3. Filtering a search using metadata filters\n",
    "    4. A similarity search over summaries of the chat messages\n",
    "    5. Using maximal marginal relevance re-ranking of a summary search\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-11T20:31:12.231459Z",
     "start_time": "2023-08-11T20:31:11.211176Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import getpass\n",
    "import time\n",
    "from uuid import uuid4\n",
    "\n",
    "from langchain.memory import ZepMemory\n",
    "from langchain_core.messages import AIMessage, HumanMessage\n",
    "\n",
    "# Set this to your Zep server URL\n",
    "ZEP_API_URL = \"http://localhost:8000\""
   ]
  },
  {
   "attachments": {},
   "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": {
    "ExecuteTime": {
     "end_time": "2023-08-11T20:31:12.237545Z",
     "start_time": "2023-08-11T20:31:12.232416Z"
    }
   },
   "outputs": [],
   "source": [
    "# Provide your Zep API key. Note that this is optional. See https://docs.getzep.com/deployment/auth\n",
    "AUTHENTICATE = False\n",
    "\n",
    "zep_api_key = None\n",
    "if AUTHENTICATE:\n",
    "    zep_api_key = getpass.getpass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-11T20:31:12.342790Z",
     "start_time": "2023-08-11T20:31:12.235291Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "session_id = str(uuid4())  # This is a unique identifier for the user/session\n",
    "\n",
    "# Initialize the Zep Memory Class\n",
    "zep_memory = ZepMemory(session_id=session_id, url=ZEP_API_URL, api_key=zep_api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-11T20:31:14.455269Z",
     "start_time": "2023-08-11T20:31:12.345635Z"
    },
    "collapsed": 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",
    "    {\"role\": \"human\", \"content\": \"What is the setting of the book?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"The book is set in a dystopian future in the 2020s, where society has\"\n",
    "            \" collapsed due to climate change and economic crises.\"\n",
    "        ),\n",
    "    },\n",
    "    {\"role\": \"human\", \"content\": \"Who is the protagonist?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"The protagonist of the book is Lauren Olamina, a young woman who possesses\"\n",
    "            \" 'hyperempathy', the ability to feel pain and other sensations she\"\n",
    "            \" witnesses.\"\n",
    "        ),\n",
    "    },\n",
    "    {\"role\": \"human\", \"content\": \"What is the main theme of the book?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"The main theme of the book is survival in the face of drastic societal\"\n",
    "            \" change and collapse. It also explores themes of adaptability, community,\"\n",
    "            \" and the human capacity for change.\"\n",
    "        ),\n",
    "    },\n",
    "    {\"role\": \"human\", \"content\": \"What is the 'Parable of the Sower'?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"The 'Parable of the Sower' is a biblical parable that Butler uses as a\"\n",
    "            \" metaphor in the book. In the parable, a sower scatters seeds, some of\"\n",
    "            \" which fall on fertile ground and grow, while others fall on rocky ground\"\n",
    "            \" or among thorns and fail to grow. The parable is used to illustrate the\"\n",
    "            \" importance of receptivity and preparedness in the face of change.\"\n",
    "        ),\n",
    "    },\n",
    "    {\"role\": \"human\", \"content\": \"What is Butler's writing style like?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"Butler's writing style is known for its clarity, directness, and\"\n",
    "            \" psychological insight. Her narratives often involve complex, diverse\"\n",
    "            \" characters and explore themes of race, gender, and power.\"\n",
    "        ),\n",
    "    },\n",
    "    {\"role\": \"human\", \"content\": \"What other books has she written?\"},\n",
    "    {\n",
    "        \"role\": \"ai\",\n",
    "        \"content\": (\n",
    "            \"In addition to 'Parable of the Sower', Butler has written several other\"\n",
    "            \" notable works, including 'Kindred', 'Dawn', and 'Parable of the Talents'.\"\n",
    "        ),\n",
    "    },\n",
    "]\n",
    "\n",
    "for msg in test_history:\n",
    "    zep_memory.chat_memory.add_message(\n",
    "        HumanMessage(content=msg[\"content\"])\n",
    "        if msg[\"role\"] == \"human\"\n",
    "        else AIMessage(content=msg[\"content\"])\n",
    "    )\n",
    "\n",
    "time.sleep(\n",
    "    10\n",
    ")  # Wait for the messages to be embedded and summarized. Speed depends on OpenAI API latency and your rate limits."
   ]
  },
  {
   "attachments": {},
   "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": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-11T20:31:14.758738Z",
     "start_time": "2023-08-11T20:31:14.458850Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content=\"What is the 'Parable of the Sower'?\", metadata={'score': 0.9250216484069824, 'uuid': '4cbfb1c0-6027-4678-af43-1e18acb224bb', 'created_at': '2023-11-01T00:32:40.224256Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 34, 'Start': 13, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}]}}, 'token_count': 13}),\n",
       " 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.8897348046302795, 'uuid': '3dd9f5ed-9dc9-4427-9da6-aba1b8278a5c', 'created_at': '2023-11-01T00:32:40.192527Z', 'updated_at': '0001-01-01T00:00:00Z', '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'}], 'intent': 'Providing information'}}, '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.8856019973754883, 'uuid': '81761dcb-38f3-4686-a4f5-6cb1007eaf29', 'created_at': '2023-11-01T00:32:40.187543Z', 'updated_at': '0001-01-01T00:00:00Z', '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'}], 'intent': \"The subject is asking for a brief summary of Butler's book, Parable of the Sower, and what it is about.\"}}, 'token_count': 23}),\n",
       " Document(page_content=\"The 'Parable of the Sower' is a biblical parable that Butler uses as a metaphor in the book. In the parable, a sower scatters seeds, some of which fall on fertile ground and grow, while others fall on rocky ground or among thorns and fail to grow. The parable is used to illustrate the importance of receptivity and preparedness in the face of change.\", metadata={'score': 0.8781436681747437, 'uuid': '1a8c5f99-2fec-425d-bc37-176ab91e7080', 'created_at': '2023-11-01T00:32:40.22836Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 26, 'Start': 5, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}, {'Label': 'ORG', 'Matches': [{'End': 60, 'Start': 54, 'Text': 'Butler'}], 'Name': 'Butler'}]}}, 'token_count': 84}),\n",
       " Document(page_content=\"In addition to 'Parable of the Sower', Butler has written several other notable works, including 'Kindred', 'Dawn', and 'Parable of the Talents'.\", metadata={'score': 0.8745182752609253, 'uuid': '45d8aa08-85ab-432f-8902-81712fe363b9', 'created_at': '2023-11-01T00:32:40.245081Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 37, 'Start': 16, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}, {'Label': 'ORG', 'Matches': [{'End': 45, 'Start': 39, 'Text': 'Butler'}], 'Name': 'Butler'}, {'Label': 'GPE', 'Matches': [{'End': 105, 'Start': 98, 'Text': 'Kindred'}], 'Name': 'Kindred'}, {'Label': 'WORK_OF_ART', 'Matches': [{'End': 144, 'Start': 121, 'Text': \"Parable of the Talents'\"}], 'Name': \"Parable of the Talents'\"}]}}, 'token_count': 39})]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.retrievers import ZepRetriever\n",
    "from langchain.retrievers.zep import SearchScope, SearchType\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?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the Zep sync API to retrieve results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-11T20:31:14.922838Z",
     "start_time": "2023-08-11T20:31:14.751737Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content=\"What is the 'Parable of the Sower'?\", metadata={'score': 0.9250596761703491, 'uuid': '4cbfb1c0-6027-4678-af43-1e18acb224bb', 'created_at': '2023-11-01T00:32:40.224256Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 34, 'Start': 13, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}]}}, 'token_count': 13}),\n",
       " 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.8897120952606201, 'uuid': '3dd9f5ed-9dc9-4427-9da6-aba1b8278a5c', 'created_at': '2023-11-01T00:32:40.192527Z', 'updated_at': '0001-01-01T00:00:00Z', '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'}], 'intent': 'Providing information'}}, '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.885666012763977, 'uuid': '81761dcb-38f3-4686-a4f5-6cb1007eaf29', 'created_at': '2023-11-01T00:32:40.187543Z', 'updated_at': '0001-01-01T00:00:00Z', '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'}], 'intent': \"The subject is asking for a brief summary of Butler's book, Parable of the Sower, and what it is about.\"}}, 'token_count': 23}),\n",
       " Document(page_content=\"The 'Parable of the Sower' is a biblical parable that Butler uses as a metaphor in the book. In the parable, a sower scatters seeds, some of which fall on fertile ground and grow, while others fall on rocky ground or among thorns and fail to grow. The parable is used to illustrate the importance of receptivity and preparedness in the face of change.\", metadata={'score': 0.878172755241394, 'uuid': '1a8c5f99-2fec-425d-bc37-176ab91e7080', 'created_at': '2023-11-01T00:32:40.22836Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 26, 'Start': 5, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}, {'Label': 'ORG', 'Matches': [{'End': 60, 'Start': 54, 'Text': 'Butler'}], 'Name': 'Butler'}]}}, 'token_count': 84}),\n",
       " Document(page_content=\"In addition to 'Parable of the Sower', Butler has written several other notable works, including 'Kindred', 'Dawn', and 'Parable of the Talents'.\", metadata={'score': 0.8745154142379761, 'uuid': '45d8aa08-85ab-432f-8902-81712fe363b9', 'created_at': '2023-11-01T00:32:40.245081Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 37, 'Start': 16, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}, {'Label': 'ORG', 'Matches': [{'End': 45, 'Start': 39, 'Text': 'Butler'}], 'Name': 'Butler'}, {'Label': 'GPE', 'Matches': [{'End': 105, 'Start': 98, 'Text': 'Kindred'}], 'Name': 'Kindred'}, {'Label': 'WORK_OF_ART', 'Matches': [{'End': 144, 'Start': 121, 'Text': \"Parable of the Talents'\"}], 'Name': \"Parable of the Talents'\"}]}}, 'token_count': 39})]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zep_retriever.get_relevant_documents(\"Who wrote Parable of the Sower?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reranking using MMR (Maximal Marginal Relevance)\n",
    "\n",
    "Zep has native, SIMD-accelerated support for reranking results using MMR. This is useful for removing redundancy in results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-11T20:31:14.923032Z",
     "start_time": "2023-08-11T20:31:14.918181Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content=\"What is the 'Parable of the Sower'?\", metadata={'score': 0.9250596761703491, 'uuid': '4cbfb1c0-6027-4678-af43-1e18acb224bb', 'created_at': '2023-11-01T00:32:40.224256Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 34, 'Start': 13, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}]}}, 'token_count': 13}),\n",
       " Document(page_content='What other books has she written?', metadata={'score': 0.77488774061203, 'uuid': '1b3c5079-9cab-46f3-beae-fb56c572e0fd', 'created_at': '2023-11-01T00:32:40.240135Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'token_count': 9}),\n",
       " Document(page_content=\"In addition to 'Parable of the Sower', Butler has written several other notable works, including 'Kindred', 'Dawn', and 'Parable of the Talents'.\", metadata={'score': 0.8745154142379761, 'uuid': '45d8aa08-85ab-432f-8902-81712fe363b9', 'created_at': '2023-11-01T00:32:40.245081Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 37, 'Start': 16, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}, {'Label': 'ORG', 'Matches': [{'End': 45, 'Start': 39, 'Text': 'Butler'}], 'Name': 'Butler'}, {'Label': 'GPE', 'Matches': [{'End': 105, 'Start': 98, 'Text': 'Kindred'}], 'Name': 'Kindred'}, {'Label': 'WORK_OF_ART', 'Matches': [{'End': 144, 'Start': 121, 'Text': \"Parable of the Talents'\"}], 'Name': \"Parable of the Talents'\"}]}}, 'token_count': 39}),\n",
       " 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.8897120952606201, 'uuid': '3dd9f5ed-9dc9-4427-9da6-aba1b8278a5c', 'created_at': '2023-11-01T00:32:40.192527Z', 'updated_at': '0001-01-01T00:00:00Z', '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'}], 'intent': 'Providing information'}}, 'token_count': 56}),\n",
       " Document(page_content='Who is the protagonist?', metadata={'score': 0.7858647704124451, 'uuid': 'ee514b37-a0b0-4d24-b0c9-3e9f8ad9d52d', 'created_at': '2023-11-01T00:32:40.203891Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'intent': 'The subject is asking about the identity of the protagonist in a specific context, such as a story, movie, or game.'}}, 'token_count': 7})]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "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",
    "    search_type=SearchType.mmr,\n",
    "    mmr_lambda=0.5,\n",
    ")\n",
    "\n",
    "await zep_retriever.aget_relevant_documents(\"Who wrote Parable of the Sower?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using metadata filters to refine search results\n",
    "\n",
    "Zep supports filtering results by metadata. This is useful for filtering results by entity type, or other metadata.\n",
    "\n",
    "More information here: https://docs.getzep.com/sdk/search_query/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content=\"What is the 'Parable of the Sower'?\", metadata={'score': 0.9251098036766052, 'uuid': '4cbfb1c0-6027-4678-af43-1e18acb224bb', 'created_at': '2023-11-01T00:32:40.224256Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 34, 'Start': 13, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}]}}, 'token_count': 13}),\n",
       " Document(page_content='What other books has she written?', metadata={'score': 0.7747920155525208, 'uuid': '1b3c5079-9cab-46f3-beae-fb56c572e0fd', 'created_at': '2023-11-01T00:32:40.240135Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'token_count': 9}),\n",
       " Document(page_content=\"In addition to 'Parable of the Sower', Butler has written several other notable works, including 'Kindred', 'Dawn', and 'Parable of the Talents'.\", metadata={'score': 0.8745266795158386, 'uuid': '45d8aa08-85ab-432f-8902-81712fe363b9', 'created_at': '2023-11-01T00:32:40.245081Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'ai', 'metadata': {'system': {'entities': [{'Label': 'WORK_OF_ART', 'Matches': [{'End': 37, 'Start': 16, 'Text': \"Parable of the Sower'\"}], 'Name': \"Parable of the Sower'\"}, {'Label': 'ORG', 'Matches': [{'End': 45, 'Start': 39, 'Text': 'Butler'}], 'Name': 'Butler'}, {'Label': 'GPE', 'Matches': [{'End': 105, 'Start': 98, 'Text': 'Kindred'}], 'Name': 'Kindred'}, {'Label': 'WORK_OF_ART', 'Matches': [{'End': 144, 'Start': 121, 'Text': \"Parable of the Talents'\"}], 'Name': \"Parable of the Talents'\"}]}}, 'token_count': 39}),\n",
       " 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.8897372484207153, 'uuid': '3dd9f5ed-9dc9-4427-9da6-aba1b8278a5c', 'created_at': '2023-11-01T00:32:40.192527Z', 'updated_at': '0001-01-01T00:00:00Z', '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'}], 'intent': 'Providing information'}}, 'token_count': 56}),\n",
       " Document(page_content='Who is the protagonist?', metadata={'score': 0.7858127355575562, 'uuid': 'ee514b37-a0b0-4d24-b0c9-3e9f8ad9d52d', 'created_at': '2023-11-01T00:32:40.203891Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'intent': 'The subject is asking about the identity of the protagonist in a specific context, such as a story, movie, or game.'}}, 'token_count': 7})]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "filter = {\"where\": {\"jsonpath\": '$[*] ? (@.Label == \"WORK_OF_ART\")'}}\n",
    "\n",
    "await zep_retriever.aget_relevant_documents(\n",
    "    \"Who wrote Parable of the Sower?\", metadata=filter\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Searching over Summaries with MMR Reranking\n",
    "\n",
    "Zep automatically generates summaries of chat messages. These summaries can be searched over using the Zep Retriever. Since a summary is a distillation of a conversation, they're more likely to match your search query and offer rich, succinct context to the LLM.\n",
    "\n",
    "Successive summaries may include similar content, with Zep's similarity search returning the highest matching results but with little diversity.\n",
    "MMR re-ranks the results to ensure that the summaries you populate into your prompt are both relevant and each offers additional information to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='The human asks about Octavia Butler and the AI informs them that she was an American science fiction author. The human\\nasks which of her books were made into movies and the AI mentions the FX series Kindred. The human then asks about her\\ncontemporaries and the AI lists Ursula K. Le Guin, Samuel R. Delany, and Joanna Russ. The human also asks about the awards\\nshe won and the AI mentions the Hugo Award, the Nebula Award, and the MacArthur Fellowship. The human asks about other women sci-fi writers to read and the AI suggests Ursula K. Le Guin and Joanna Russ. The human then asks for a synopsis of Butler\\'s book \"Parable of the Sower\" and the AI describes it.', metadata={'score': 0.7882999777793884, 'uuid': '3c95a29a-52dc-4112-b8a7-e6b1dc414d45', 'created_at': '2023-11-01T00:32:47.76449Z', 'token_count': 155}),\n",
       " Document(page_content='The human asks about Octavia Butler. The AI informs the human that Octavia Estelle Butler was an American science \\nfiction author. The human then asks which books of hers were made into movies and the AI mentions the FX series Kindred, \\nbased on her novel of the same name.', metadata={'score': 0.7407922744750977, 'uuid': '0e027f4d-d71f-42ae-977f-696b8948b8bf', 'created_at': '2023-11-01T00:32:41.637098Z', 'token_count': 59}),\n",
       " Document(page_content='The human asks about Octavia Butler and the AI informs them that she was an American science fiction author. The human\\nasks which of her books were made into movies and the AI mentions the FX series Kindred. The human then asks about her\\ncontemporaries and the AI lists Ursula K. Le Guin, Samuel R. Delany, and Joanna Russ. The human also asks about the awards\\nshe won and the AI mentions the Hugo Award, the Nebula Award, and the MacArthur Fellowship.', metadata={'score': 0.7436535358428955, 'uuid': 'b3500d1b-1a78-4aef-9e24-6b196cfa83cb', 'created_at': '2023-11-01T00:32:44.24744Z', 'token_count': 104})]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "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=3,\n",
    "    api_key=zep_api_key,\n",
    "    search_scope=SearchScope.summary,\n",
    "    search_type=SearchType.mmr,\n",
    "    mmr_lambda=0.5,\n",
    ")\n",
    "\n",
    "await zep_retriever.aget_relevant_documents(\"Who wrote Parable of the Sower?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}