From 20e7d8c4630518159b436c4d5c4f3a898de382ef Mon Sep 17 00:00:00 2001 From: Bagatur Date: Fri, 11 Aug 2023 10:36:08 -0700 Subject: [PATCH] rfc --- libs/langchain/langchain/memory/zep_memory.py | 121 +++++++++++------- libs/langchain/langchain/retrievers/zep.py | 2 +- 2 files changed, 77 insertions(+), 46 deletions(-) diff --git a/libs/langchain/langchain/memory/zep_memory.py b/libs/langchain/langchain/memory/zep_memory.py index 6e2a7ec823b..f0f889a6397 100644 --- a/libs/langchain/langchain/memory/zep_memory.py +++ b/libs/langchain/langchain/memory/zep_memory.py @@ -1,56 +1,16 @@ from __future__ import annotations -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional from langchain.memory import ConversationBufferMemory +from langchain.memory.chat_memory import BaseChatMemory from langchain.memory.chat_message_histories import ZepChatMessageHistory +from langchain.memory.utils import get_prompt_input_key -class ZepMemory(ConversationBufferMemory): - """Persist your chain history to the Zep Memory Server. - - The number of messages returned by Zep and when the Zep server summarizes chat - histories is configurable. See the Zep documentation for more details. - - Documentation: https://docs.getzep.com - - Example: - .. code-block:: python - - memory = ZepMemory( - session_id=session_id, # Identifies your user or a user's session - url=ZEP_API_URL, # Your Zep server's URL - api_key=, # Optional - memory_key="history", # Ensure this matches the key used in - # chain's prompt template - return_messages=True, # Does your prompt template expect a string - # or a list of Messages? - ) - chain = LLMChain(memory=memory,...) # Configure your chain to use the ZepMemory - instance - - - Note: - To persist metadata alongside your chat history, your will need to create a - custom Chain class that overrides the `prep_outputs` method to include the metadata - in the call to `self.memory.save_context`. - - - About 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 and more, see: - https://docs.getzep.com/deployment/quickstart/ - - For more information on the zep-python package, see: - https://github.com/getzep/zep-python - - """ - +class _ZepMemory(BaseChatMemory): chat_memory: ZepChatMessageHistory + memory_key: str = "history" #: :meta private: def __init__( self, @@ -102,6 +62,77 @@ class ZepMemory(ConversationBufferMemory): memory_key=memory_key, ) + @property + def memory_variables(self) -> List[str]: + """Will always return list of memory variables. + + :meta private: + """ + return [self.memory_key] + + +class ZepSearchMemory(_ZepMemory): + top_k: int = 4 + + def _get_prompt_input_key(self, inputs: Dict[str, Any]) -> str: + """Get the input key for the prompt.""" + if self.input_key is None: + return get_prompt_input_key(inputs, self.memory_variables) + return self.input_key + + def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, str]: + """Return history buffer.""" + input_key = self._get_prompt_input_key(inputs) + query = inputs[input_key] + results = self.chat_memory.search(query, limit=self.top_k) + result = "\n".join([r.message.pop("content") for r in results]) + return {self.memory_key: result} + + +class ZepBufferMemory(_ZepMemory, ConversationBufferMemory): + """Persist your chain history to the Zep Memory Server. + + The number of messages returned by Zep and when the Zep server summarizes chat + histories is configurable. See the Zep documentation for more details. + + Documentation: https://docs.getzep.com + + Example: + .. code-block:: python + + memory = ZepMemory( + session_id=session_id, # Identifies your user or a user's session + url=ZEP_API_URL, # Your Zep server's URL + api_key=, # Optional + memory_key="history", # Ensure this matches the key used in + # chain's prompt template + return_messages=True, # Does your prompt template expect a string + # or a list of Messages? + ) + chain = LLMChain(memory=memory,...) # Configure your chain to use the ZepMemory + instance + + + Note: + To persist metadata alongside your chat history, your will need to create a + custom Chain class that overrides the `prep_outputs` method to include the metadata + in the call to `self.memory.save_context`. + + + About 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 and more, see: + https://docs.getzep.com/deployment/quickstart/ + + For more information on the zep-python package, see: + https://github.com/getzep/zep-python + + """ + def save_context( self, inputs: Dict[str, Any], diff --git a/libs/langchain/langchain/retrievers/zep.py b/libs/langchain/langchain/retrievers/zep.py index 7d1f18cda1a..677640cb7c1 100644 --- a/libs/langchain/langchain/retrievers/zep.py +++ b/libs/langchain/langchain/retrievers/zep.py @@ -42,7 +42,7 @@ class ZepRetriever(BaseRetriever): try: from zep_python import ZepClient except ImportError: - raise ValueError( + raise ImportError( "Could not import zep-python package. " "Please install it with `pip install zep-python`." )