From f6601fe5ffc60bc9992450b4da83a0dda18dc101 Mon Sep 17 00:00:00 2001 From: Daniel Chalef <131175+danielchalef@users.noreply.github.com> Date: Tue, 31 Oct 2023 17:34:29 -0700 Subject: [PATCH] add summary search and example --- .../retrievers/zep_memorystore.ipynb | 178 +++++++++++++----- libs/langchain/langchain/retrievers/zep.py | 52 ++++- 2 files changed, 177 insertions(+), 53 deletions(-) diff --git a/docs/docs/integrations/retrievers/zep_memorystore.ipynb b/docs/docs/integrations/retrievers/zep_memorystore.ipynb index 2f91fd2b0ffd61..eda3c6a04b05ce 100644 --- a/docs/docs/integrations/retrievers/zep_memorystore.ipynb +++ b/docs/docs/integrations/retrievers/zep_memorystore.ipynb @@ -18,7 +18,7 @@ "\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", + "- **Auto-summarization** of memory messages based on a configurable message window. A series of summaries are stored, and can be searched and retrieved.\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", @@ -40,7 +40,12 @@ "We'll demonstrate:\n", "\n", "1. Adding conversation history to the Zep memory store.\n", - "2. Vector search over the conversation history.\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" ] }, @@ -106,16 +111,7 @@ }, "collapsed": false }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/danielchalef/dev/langchain/.venv/lib/python3.11/site-packages/zep_python/zep_client.py:86: Warning: You are using an incompatible Zep server version. Please upgrade to {MINIMUM_SERVER_VERSION} or later.\n", - " self._healthcheck(base_url)\n" - ] - } - ], + "outputs": [], "source": [ "session_id = str(uuid4()) # This is a unique identifier for the user/session\n", "\n", @@ -193,6 +189,60 @@ " \" 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", @@ -202,7 +252,9 @@ " else AIMessage(content=msg[\"content\"])\n", " )\n", "\n", - "time.sleep(2) # Wait for the messages to be embedded" + "time.sleep(\n", + " 10\n", + ") # Wait for the messages to be embedded and summarized. Speed depends on OpenAI API latency and your rate limits." ] }, { @@ -231,11 +283,11 @@ { "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.8897589445114136, 'uuid': 'f99ecec3-f778-4bfd-8bb7-c3c00ae919c0', 'created_at': '2023-10-17T22:53:08.664849Z', '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'}]}}, '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.8856973648071289, 'uuid': 'f6aba470-f15f-4b22-84ef-1c0d315a31de', 'created_at': '2023-10-17T22:53:08.642659Z', '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'}]}}, 'token_count': 23}),\n", - " Document(page_content='Who was Octavia Butler?', metadata={'score': 0.7759557962417603, 'uuid': '26aab7b5-34b1-4aff-9be0-7834a7702be4', 'created_at': '2023-10-17T22:53:08.585297Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'PERSON', 'Matches': [{'End': 22, 'Start': 8, 'Text': 'Octavia Butler'}], 'Name': 'Octavia Butler'}], 'intent': 'The subject is asking for information about Octavia Butler, a specific person.'}}, '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.760245680809021, 'uuid': 'ee4aa8e9-9913-4e69-a2a5-77a85294d24e', 'created_at': '2023-10-17T22:53:08.611466Z', 'updated_at': '0001-01-01T00:00:00Z', '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'}]}}, 'token_count': 27}),\n", - " Document(page_content='You might want to read Ursula K. Le Guin or Joanna Russ.', metadata={'score': 0.7596070170402527, 'uuid': '9fa630e6-0b17-4d77-80b0-ba99249850c0', 'created_at': '2023-10-17T22:53:08.630731Z', 'updated_at': '0001-01-01T00:00:00Z', '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})]" + "[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, @@ -245,7 +297,7 @@ ], "source": [ "from langchain.retrievers import ZepRetriever\n", - "from langchain.retrievers.zep import SearchType\n", + "from langchain.retrievers.zep import SearchType, SearchScope\n", "\n", "zep_retriever = ZepRetriever(\n", " session_id=session_id, # Ensure that you provide the session_id when instantiating the Retriever\n", @@ -279,11 +331,11 @@ { "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.8897120952606201, 'uuid': 'f99ecec3-f778-4bfd-8bb7-c3c00ae919c0', 'created_at': '2023-10-17T22:53:08.664849Z', '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'}]}}, '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.8857351541519165, 'uuid': 'f6aba470-f15f-4b22-84ef-1c0d315a31de', 'created_at': '2023-10-17T22:53:08.642659Z', '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'}]}}, 'token_count': 23}),\n", - " Document(page_content='Who was Octavia Butler?', metadata={'score': 0.7759560942649841, 'uuid': '26aab7b5-34b1-4aff-9be0-7834a7702be4', 'created_at': '2023-10-17T22:53:08.585297Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'entities': [{'Label': 'PERSON', 'Matches': [{'End': 22, 'Start': 8, 'Text': 'Octavia Butler'}], 'Name': 'Octavia Butler'}], 'intent': 'The subject is asking for information about Octavia Butler, a specific person.'}}, '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.7602507472038269, 'uuid': 'ee4aa8e9-9913-4e69-a2a5-77a85294d24e', 'created_at': '2023-10-17T22:53:08.611466Z', 'updated_at': '0001-01-01T00:00:00Z', '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 stating a fact about Octavia Butler's contemporaries, including Ursula K. Le Guin, Samuel R. Delany, and Joanna Russ.\"}}, 'token_count': 27}),\n", - " Document(page_content='You might want to read Ursula K. Le Guin or Joanna Russ.', metadata={'score': 0.7595934867858887, 'uuid': '9fa630e6-0b17-4d77-80b0-ba99249850c0', 'created_at': '2023-10-17T22:53:08.630731Z', 'updated_at': '0001-01-01T00:00:00Z', '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})]" + "[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, @@ -316,22 +368,14 @@ "collapsed": false }, "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/danielchalef/dev/langchain/.venv/lib/python3.11/site-packages/zep_python/zep_client.py:86: Warning: You are using an incompatible Zep server version. Please upgrade to {MINIMUM_SERVER_VERSION} or later.\n", - " self._healthcheck(base_url)\n" - ] - }, { "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.8897120952606201, 'uuid': 'f99ecec3-f778-4bfd-8bb7-c3c00ae919c0', 'created_at': '2023-10-17T22:53:08.664849Z', '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'}]}}, 'token_count': 56}),\n", - " Document(page_content='Which books of hers were made into movies?', metadata={'score': 0.7496200799942017, 'uuid': '1047ff15-96f1-4101-bb0f-9ed073b8081d', 'created_at': '2023-10-17T22:53:08.596614Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'intent': 'The subject is inquiring about the books of the person referred to as \"hers\" that have been made into movies.'}}, 'token_count': 11}),\n", - " Document(page_content=\"Write a short synopsis of Butler's book, Parable of the Sower. What is it about?\", metadata={'score': 0.8857351541519165, 'uuid': 'f6aba470-f15f-4b22-84ef-1c0d315a31de', 'created_at': '2023-10-17T22:53:08.642659Z', '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'}]}}, 'token_count': 23}),\n", - " Document(page_content='You might want to read Ursula K. Le Guin or Joanna Russ.', metadata={'score': 0.7595934867858887, 'uuid': '9fa630e6-0b17-4d77-80b0-ba99249850c0', 'created_at': '2023-10-17T22:53:08.630731Z', 'updated_at': '0001-01-01T00:00:00Z', '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}),\n", - " Document(page_content='Who were her contemporaries?', metadata={'score': 0.7575579881668091, 'uuid': 'b2dfd1f7-cac6-4e37-94ea-7c15b0a5af2c', 'created_at': '2023-10-17T22:53:08.606283Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'intent': 'The subject is asking about the people who were contemporaries of someone else.'}}, 'token_count': 8})]" + "[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, @@ -366,20 +410,20 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 8, "metadata": {}, "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.8897120952606201, 'uuid': 'f99ecec3-f778-4bfd-8bb7-c3c00ae919c0', 'created_at': '2023-10-17T22:53:08.664849Z', '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': 'None'}}, 'token_count': 56}),\n", - " Document(page_content='Which books of hers were made into movies?', metadata={'score': 0.7496200799942017, 'uuid': '1047ff15-96f1-4101-bb0f-9ed073b8081d', 'created_at': '2023-10-17T22:53:08.596614Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'intent': 'The subject is inquiring about the books of the person referred to as \"hers\" that have been made into movies.'}}, 'token_count': 11}),\n", - " Document(page_content=\"Write a short synopsis of Butler's book, Parable of the Sower. What is it about?\", metadata={'score': 0.8857351541519165, 'uuid': 'f6aba470-f15f-4b22-84ef-1c0d315a31de', 'created_at': '2023-10-17T22:53:08.642659Z', '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 requesting a brief summary or description of Butler\\'s book, \"Parable of the Sower.\"'}}, 'token_count': 23}),\n", - " Document(page_content='You might want to read Ursula K. Le Guin or Joanna Russ.', metadata={'score': 0.7595934867858887, 'uuid': '9fa630e6-0b17-4d77-80b0-ba99249850c0', 'created_at': '2023-10-17T22:53:08.630731Z', 'updated_at': '0001-01-01T00:00:00Z', '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'}], 'intent': 'The subject is providing a suggestion or recommendation for the person to read Ursula K. Le Guin or Joanna Russ.'}}, 'token_count': 18}),\n", - " Document(page_content='Who were her contemporaries?', metadata={'score': 0.7575579881668091, 'uuid': 'b2dfd1f7-cac6-4e37-94ea-7c15b0a5af2c', 'created_at': '2023-10-17T22:53:08.606283Z', 'updated_at': '0001-01-01T00:00:00Z', 'role': 'human', 'metadata': {'system': {'intent': 'The subject is asking about the people who were contemporaries of someone else.'}}, 'token_count': 8})]" + "[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": 9, + "execution_count": 8, "metadata": {}, "output_type": "execute_result" } @@ -392,6 +436,50 @@ ")" ] }, + { + "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, @@ -416,7 +504,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.5" + "version": "3.11.6" } }, "nbformat": 4, diff --git a/libs/langchain/langchain/retrievers/zep.py b/libs/langchain/langchain/retrievers/zep.py index ae53aaf2546406..8253a35082dec8 100644 --- a/libs/langchain/langchain/retrievers/zep.py +++ b/libs/langchain/langchain/retrievers/zep.py @@ -11,7 +11,16 @@ from langchain.schema import BaseRetriever, Document if TYPE_CHECKING: - from zep_python import MemorySearchResult + from zep_python.memory import MemorySearchResult + + +class SearchScope(str, Enum): + """Which documents to search. Messages or Summaries?""" + + messages = "messages" + """Search chat history messages.""" + summary = "summary" + """Search chat history summaries.""" class SearchType(str, Enum): @@ -62,6 +71,8 @@ class ZepRetriever(BaseRetriever): """Zep session ID.""" top_k: Optional[int] """Number of items to return.""" + search_scope: SearchScope = SearchScope.messages + """Which documents to search. Messages or Summaries?""" search_type: SearchType = SearchType.similarity """Type of search to perform (similarity / mmr)""" mmr_lambda: Optional[float] = None @@ -82,7 +93,7 @@ def create_client(cls, values: dict) -> dict: ) return values - def _search_result_to_doc( + def _messages_search_result_to_doc( self, results: List[MemorySearchResult] ) -> List[Document]: return [ @@ -94,6 +105,23 @@ def _search_result_to_doc( if r.message ] + def _summary_search_result_to_doc( + self, results: List[MemorySearchResult] + ) -> List[Document]: + return [ + Document( + page_content=r.summary.content, + metadata={ + "score": r.dist, + "uuid": r.summary.uuid, + "created_at": r.summary.created_at, + "token_count": r.summary.token_count, + }, + ) + for r in results + if r.summary + ] + def _get_relevant_documents( self, query: str, @@ -101,14 +129,15 @@ def _get_relevant_documents( run_manager: CallbackManagerForRetrieverRun, metadata: Optional[Dict[str, Any]] = None, ) -> List[Document]: - from zep_python import MemorySearchPayload + from zep_python.memory import MemorySearchPayload if not self.zep_client: raise RuntimeError("Zep client not initialized.") - payload: MemorySearchPayload = MemorySearchPayload( + payload = MemorySearchPayload( text=query, metadata=metadata, + search_scope=self.search_scope, search_type=self.search_type, mmr_lambda=self.mmr_lambda, ) @@ -117,7 +146,10 @@ def _get_relevant_documents( self.session_id, payload, limit=self.top_k ) - return self._search_result_to_doc(results) + if self.search_scope == SearchScope.summary: + return self._summary_search_result_to_doc(results) + + return self._messages_search_result_to_doc(results) async def _aget_relevant_documents( self, @@ -126,14 +158,15 @@ async def _aget_relevant_documents( run_manager: AsyncCallbackManagerForRetrieverRun, metadata: Optional[Dict[str, Any]] = None, ) -> List[Document]: - from zep_python import MemorySearchPayload + from zep_python.memory import MemorySearchPayload if not self.zep_client: raise RuntimeError("Zep client not initialized.") - payload: MemorySearchPayload = MemorySearchPayload( + payload = MemorySearchPayload( text=query, metadata=metadata, + search_scope=self.search_scope, search_type=self.search_type, mmr_lambda=self.mmr_lambda, ) @@ -142,4 +175,7 @@ async def _aget_relevant_documents( self.session_id, payload, limit=self.top_k ) - return self._search_result_to_doc(results) + if self.search_scope == SearchScope.summary: + return self._summary_search_result_to_doc(results) + + return self._messages_search_result_to_doc(results)