Skip to content

Commit

Permalink
zep retriever and notebook
Browse files Browse the repository at this point in the history
  • Loading branch information
Daniel Chalef committed May 11, 2023
1 parent 446b60d commit 39d2854
Show file tree
Hide file tree
Showing 3 changed files with 370 additions and 0 deletions.
287 changes: 287 additions & 0 deletions docs/modules/indexes/retrievers/examples/zep_memorystore.ipynb
View file
Original file line number Diff line number Diff line change
@@ -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
}
2 changes: 2 additions & 0 deletions langchain/retrievers/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -34,4 +35,5 @@
"WeaviateHybridSearchRetriever",
"AzureCognitiveSearchRetriever",
"WikipediaRetriever",
"ZepRetriever",
]
81 changes: 81 additions & 0 deletions langchain/retrievers/zep.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
from __future__ import annotations

from typing import List, Optional

from langchain.schema import BaseRetriever, Document


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
"""

url: str
session_id: str
top_k: Optional[int]

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.url = url
self.top_k = top_k

def _search_result_to_doc(self, results) -> 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.models 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.models 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)

0 comments on commit 39d2854

Please sign in to comment.