Zep Retriever - Vector Search Over Chat History

docs/modules/indexes/retrievers/examples/zep_memorystore.ipynb

@@ -0,0 +1,287 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Zep Long-term Memory Store\n",
+    "\n",
+    "## Chat Message History Retriever Example\n",
+    "\n",
+    "This notebook demonstrates how to search historical chat message histories using the [Zep Long-term Memory Store](https://github.com/getzep/zep).\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",
+    "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",
+    "- 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",
+    "- Vector search over memories, with messages automatically embedded on creation.\n",
+    "- Auto-token counting of memories and summaries, allowing finer-grained control over prompt assembly.\n",
+    "- Python and JavaScript SDKs.\n",
+    "\n",
+    "Zep's Go Extractor model is easily extensible, with a simple, clean interface available to build new enrichment functionality, such as summarizers, entity extractors, embedders, and more.\n",
+    "\n",
+    "Zep project: [https://github.com/getzep/zep](https://github.com/getzep/zep)\n"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "outputs": [],
+   "source": [
+    "from langchain.memory.chat_message_histories import ZepChatMessageHistory\n",
+    "from langchain.schema import HumanMessage, AIMessage\n",
+    "from uuid import uuid4\n",
+    "\n",
+    "# Set this to your Zep server URL\n",
+    "ZEP_API_URL = \"http://localhost:8000\"\n",
+    "\n",
+    "# Zep is async-first. Our sync APIs use an asyncio wrapper to run outside an app's event loop.\n",
+    "# This interferes with Jupyter's event loop, so we need to install nest_asyncio to run the\n",
+    "# Zep client in a notebook.\n",
+    "\n",
+    "# !pip install nest_asyncio  # Uncomment to install nest_asyncio\n",
+    "import nest_asyncio\n",
+    "\n",
+    "nest_asyncio.apply()"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2023-05-11T16:29:34.865270Z",
+     "start_time": "2023-05-11T16:29:34.139039Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "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."
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "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,\n",
+    "    url=ZEP_API_URL,\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2023-05-11T16:29:34.943646Z",
+     "start_time": "2023-05-11T16:29:34.865819Z"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "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.append(\n",
+    "        HumanMessage(content=msg[\"content\"])\n",
+    "        if msg[\"role\"] == \"human\"\n",
+    "        else AIMessage(content=msg[\"content\"])\n",
+    "    )\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2023-05-11T16:29:35.083074Z",
+     "start_time": "2023-05-11T16:29:34.943496Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "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."
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "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={'source': 'd2d364f3-d2cc-46be-9e04-b41333f80514', 'score': 0.8897276230611862, 'role': 'ai', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.086398Z'}),\n Document(page_content=\"Write a short synopsis of Butler's book, Parable of the Sower. What is it about?\", metadata={'source': '2baee3f7-cfc8-4cce-9e9f-6e448249244b', 'score': 0.8858500630231024, 'role': 'human', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.083828Z'}),\n Document(page_content='Who was Octavia Butler?', metadata={'source': '73fb8cb3-da77-4f99-ac2b-939d5019e5b8', 'score': 0.7759728831215438, 'role': 'human', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.014421Z'}),\n Document(page_content=\"Octavia Butler's contemporaries included Ursula K. Le Guin, Samuel R. Delany, and Joanna Russ.\", metadata={'source': 'aff1b45d-1e14-427d-a5a2-2b5a9dade294', 'score': 0.760286350496536, 'role': 'ai', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.052896Z'}),\n Document(page_content='You might want to read Ursula K. Le Guin or Joanna Russ.', metadata={'source': '0dd8cde5-860e-4d8b-975f-50f55028177d', 'score': 0.7595191167162665, 'role': 'ai', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.080817Z'})]"
+     },
+     "execution_count": 9,
+     "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",
+    ")\n",
+    "\n",
+    "await zep_retriever.aget_relevant_documents(\"Who wrote Parable of the Sower?\")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2023-05-11T16:32:59.482138Z",
+     "start_time": "2023-05-11T16:32:59.256622Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "We can also use the Zep sync API to retrieve results:"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "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={'source': 'd2d364f3-d2cc-46be-9e04-b41333f80514', 'score': 0.8897364040715323, 'role': 'ai', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.086398Z'}),\n Document(page_content=\"Write a short synopsis of Butler's book, Parable of the Sower. What is it about?\", metadata={'source': '2baee3f7-cfc8-4cce-9e9f-6e448249244b', 'score': 0.8857559512124905, 'role': 'human', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.083828Z'}),\n Document(page_content='Who was Octavia Butler?', metadata={'source': '73fb8cb3-da77-4f99-ac2b-939d5019e5b8', 'score': 0.7759001673780126, 'role': 'human', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.014421Z'}),\n Document(page_content=\"Octavia Butler's contemporaries included Ursula K. Le Guin, Samuel R. Delany, and Joanna Russ.\", metadata={'source': 'aff1b45d-1e14-427d-a5a2-2b5a9dade294', 'score': 0.7602262941130749, 'role': 'ai', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.052896Z'}),\n Document(page_content='You might want to read Ursula K. Le Guin or Joanna Russ.', metadata={'source': '0dd8cde5-860e-4d8b-975f-50f55028177d', 'score': 0.7594491870036507, 'role': 'ai', 'token_count': 0, 'created_at': '2023-05-11T16:29:35.080817Z'})]"
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "zep_retriever.get_relevant_documents(\"Who wrote Parable of the Sower?\")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2023-05-11T16:38:44.372773Z",
+     "start_time": "2023-05-11T16:38:43.946464Z"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

langchain/retrievers/__init__.py

@@ -16,6 +16,7 @@
 from langchain.retrievers.vespa_retriever import VespaRetriever
 from langchain.retrievers.weaviate_hybrid_search import WeaviateHybridSearchRetriever
 from langchain.retrievers.wikipedia import WikipediaRetriever
+from langchain.retrievers.zep import ZepRetriever
 
 __all__ = [
     "ChatGPTPluginRetriever",
@@ -34,4 +35,5 @@
     "WeaviateHybridSearchRetriever",
     "AzureCognitiveSearchRetriever",
     "WikipediaRetriever",
+    "ZepRetriever",
 ]

langchain/retrievers/zep.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
+
+from langchain.schema import BaseRetriever, Document
+
+if TYPE_CHECKING:
+    from zep_python import SearchResult
+
+
+class ZepRetriever(BaseRetriever):
+    """A Retriever implementation for the Zep long-term memory store. Search your
+    user's long-term chat history with Zep.
+
+    Note: You will need to provide the user's `session_id` to use this retriever.
+
+    More on Zep:
+    Zep provides long-term conversation storage for LLM apps. The server stores,
+    summarizes, embeds, indexes, and enriches conversational AI chat
+    histories, and exposes them via simple, low-latency APIs.
+
+    For server installation instructions, see: https://github.com/getzep/zep
+    """
+
+    def __init__(
+        self,
+        session_id: str,
+        url: str,
+        top_k: Optional[int] = None,
+    ):
+        try:
+            from zep_python import ZepClient
+        except ImportError:
+            raise ValueError(
+                "Could not import zep-python package. "
+                "Please install it with `pip install zep-python`."
+            )
+
+        self.zep_client = ZepClient(base_url=url)
+        self.session_id = session_id
+        self.top_k = top_k
+
+    def _search_result_to_doc(self, results: List[SearchResult]) -> List[Document]:
+        return [
+            Document(
+                page_content=r.message["content"],
+                metadata={
+                    "source": r.message["uuid"],
+                    "score": r.dist,
+                    "role": r.message["role"],
+                    "token_count": r.message["token_count"],
+                    "created_at": r.message["created_at"],
+                },
+            )
+            for r in results
+            if r.message
+        ]
+
+    def get_relevant_documents(self, query: str) -> List[Document]:
+        from zep_python import SearchPayload, SearchResult
+
+        payload: SearchPayload = SearchPayload(text=query)
+
+        results: List[SearchResult] = self.zep_client.search_memory(
+            self.session_id, payload, limit=self.top_k
+        )
+
+        return self._search_result_to_doc(results)
+
+    async def aget_relevant_documents(self, query: str) -> List[Document]:
+        from zep_python import SearchPayload, SearchResult
+
+        payload: SearchPayload = SearchPayload(text=query)
+
+        results: List[SearchResult] = await self.zep_client.asearch_memory(
+            self.session_id, payload, limit=self.top_k
+        )
+
+        return self._search_result_to_doc(results)