Source code for langchain.memory.vectorstore_token_buffer_memory

"""
Class for a conversation memory buffer with older messages stored in a vectorstore .

This implements a conversation memory in which the messages are stored in a memory
buffer up to a specified token limit. When the limit is exceeded, older messages are
saved to a vectorstore backing database. The vectorstore can be made persistent across
sessions.
"""

import warnings
from datetime import datetime
from typing import Any, Dict, List

from langchain_core.messages import BaseMessage
from langchain_core.prompts.chat import SystemMessagePromptTemplate
from langchain_core.pydantic_v1 import Field, PrivateAttr
from langchain_core.vectorstores import VectorStoreRetriever

from langchain.memory import ConversationTokenBufferMemory, VectorStoreRetrieverMemory
from langchain.memory.chat_memory import BaseChatMemory
from langchain.text_splitter import RecursiveCharacterTextSplitter

DEFAULT_HISTORY_TEMPLATE = """
Current date and time: {current_time}.

Potentially relevant timestamped excerpts of previous conversations (you 
do not need to use these if irrelevant):
{previous_history}

"""

TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S %Z"


[docs]class ConversationVectorStoreTokenBufferMemory(ConversationTokenBufferMemory): """Conversation chat memory with token limit and vectordb backing. load_memory_variables() will return a dict with the key "history". It contains background information retrieved from the vector store plus recent lines of the current conversation. To help the LLM understand the part of the conversation stored in the vectorstore, each interaction is timestamped and the current date and time is also provided in the history. A side effect of this is that the LLM will have access to the current date and time. Initialization arguments: This class accepts all the initialization arguments of ConversationTokenBufferMemory, such as `llm`. In addition, it accepts the following additional arguments retriever: (required) A VectorStoreRetriever object to use as the vector backing store split_chunk_size: (optional, 1000) Token chunk split size for long messages generated by the AI previous_history_template: (optional) Template used to format the contents of the prompt history Example using ChromaDB: .. code-block:: python from langchain.memory.token_buffer_vectorstore_memory import ( ConversationVectorStoreTokenBufferMemory ) from langchain_chroma import Chroma from langchain_community.embeddings import HuggingFaceInstructEmbeddings from langchain_openai import OpenAI embedder = HuggingFaceInstructEmbeddings( query_instruction="Represent the query for retrieval: " ) chroma = Chroma(collection_name="demo", embedding_function=embedder, collection_metadata={"hnsw:space": "cosine"}, ) retriever = chroma.as_retriever( search_type="similarity_score_threshold", search_kwargs={ 'k': 5, 'score_threshold': 0.75, }, ) conversation_memory = ConversationVectorStoreTokenBufferMemory( return_messages=True, llm=OpenAI(), retriever=retriever, max_token_limit = 1000, ) conversation_memory.save_context({"Human": "Hi there"}, {"AI": "Nice to meet you!"} ) conversation_memory.save_context({"Human": "Nice day isn't it?"}, {"AI": "I love Wednesdays."} ) conversation_memory.load_memory_variables({"input": "What time is it?"}) """ retriever: VectorStoreRetriever = Field(exclude=True) memory_key: str = "history" previous_history_template: str = DEFAULT_HISTORY_TEMPLATE split_chunk_size: int = 1000 _memory_retriever: VectorStoreRetrieverMemory = PrivateAttr(default=None) _timestamps: List[datetime] = PrivateAttr(default_factory=list) @property def memory_retriever(self) -> VectorStoreRetrieverMemory: """Return a memory retriever from the passed retriever object.""" if self._memory_retriever is not None: return self._memory_retriever self._memory_retriever = VectorStoreRetrieverMemory(retriever=self.retriever) return self._memory_retriever
[docs] def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]: """Return history and memory buffer.""" try: with warnings.catch_warnings(): warnings.simplefilter("ignore") memory_variables = self.memory_retriever.load_memory_variables(inputs) previous_history = memory_variables[self.memory_retriever.memory_key] except AssertionError: # happens when db is empty previous_history = "" current_history = super().load_memory_variables(inputs) template = SystemMessagePromptTemplate.from_template( self.previous_history_template ) messages = [ template.format( previous_history=previous_history, current_time=datetime.now().astimezone().strftime(TIMESTAMP_FORMAT), ) ] messages.extend(current_history[self.memory_key]) return {self.memory_key: messages}
[docs] def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None: """Save context from this conversation to buffer. Pruned.""" BaseChatMemory.save_context(self, inputs, outputs) self._timestamps.append(datetime.now().astimezone()) # Prune buffer if it exceeds max token limit buffer = self.chat_memory.messages curr_buffer_length = self.llm.get_num_tokens_from_messages(buffer) if curr_buffer_length > self.max_token_limit: while curr_buffer_length > self.max_token_limit: self._pop_and_store_interaction(buffer) curr_buffer_length = self.llm.get_num_tokens_from_messages(buffer)
[docs] def save_remainder(self) -> None: """ Save the remainder of the conversation buffer to the vector store. This is useful if you have made the vectorstore persistent, in which case this can be called before the end of the session to store the remainder of the conversation. """ buffer = self.chat_memory.messages while len(buffer) > 0: self._pop_and_store_interaction(buffer)
def _pop_and_store_interaction(self, buffer: List[BaseMessage]) -> None: input = buffer.pop(0) output = buffer.pop(0) timestamp = self._timestamps.pop(0).strftime(TIMESTAMP_FORMAT) # Split AI output into smaller chunks to avoid creating documents # that will overflow the context window ai_chunks = self._split_long_ai_text(str(output.content)) for index, chunk in enumerate(ai_chunks): self.memory_retriever.save_context( {"Human": f"<{timestamp}/00> {str(input.content)}"}, {"AI": f"<{timestamp}/{index:02}> {chunk}"}, ) def _split_long_ai_text(self, text: str) -> List[str]: splitter = RecursiveCharacterTextSplitter(chunk_size=self.split_chunk_size) return [chunk.page_content for chunk in splitter.create_documents([text])]