Repository: tmc/langchaingo
Pull request #222 (2023-07-25 15:06:42)

memory: Feature conversational retrieval interface fix and rename
### PR Checklist

- [x] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [x] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [x] Name your Pull Request title clearly, concisely, and prefixed with the name of the primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [x] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [x] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `Fixes #123`).
- [x] Describes the source of new concepts.
- [x] References existing implementations as appropriate.
- [x] Contains test coverage for new functions.
- [x] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.

## Description
- this PR fixes the Conversational retrieval interface to accept `schema.Memory` instead of the `memory.Buffer`, by that we can in future pass in any memory that we need (for example Buffer, TokenBuffer or any other that we add).
- I've also renamed `Buffer` to be `ConversationBuffer` and `TokenBuffer` to be `ConversationTokenBuffer` so that we are closer to the langchain python implementation, so that it is easier on the eyes when we are comparing what is what.
- updated `ChatMessageHistory` schema to support errors, we need that when we port in psql, mongo or any other driven history, so that we don't swallow the errors when/if we have them


Patch:
diff --git a/chains/conversational_retrieval_qa.go b/chains/conversational_retrieval_qa.go
index 2d58a6ce..f154fc37 100644
--- a/chains/conversational_retrieval_qa.go
+++ b/chains/conversational_retrieval_qa.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 
 	"github.com/tmc/langchaingo/llms"
-	"github.com/tmc/langchaingo/memory"
 	"github.com/tmc/langchaingo/schema"
 )
 
@@ -20,8 +19,8 @@ type ConversationalRetrievalQA struct {
 	// Retriever used to retrieve the relevant documents.
 	Retriever schema.Retriever
 
-	// Buffer is a simple form of memory that remembers previous conversational back and forths directly.
-	Memory memory.Buffer
+	// Memory that remembers previous conversational back and forths directly.
+	Memory schema.Memory
 
 	// CombineDocumentsChain The chain used to combine any retrieved documents.
 	CombineDocumentsChain Chain
@@ -59,7 +58,7 @@ func NewConversationalRetrievalQA(
 	combineDocumentsChain Chain,
 	condenseQuestionChain Chain,
 	retriever schema.Retriever,
-	memory memory.Buffer,
+	memory schema.Memory,
 ) ConversationalRetrievalQA {
 	return ConversationalRetrievalQA{
 		Memory:                  memory,
@@ -77,7 +76,7 @@ func NewConversationalRetrievalQA(
 func NewConversationalRetrievalQAFromLLM(
 	llm llms.LanguageModel,
 	retriever schema.Retriever,
-	memory memory.Buffer,
+	memory schema.Memory,
 ) ConversationalRetrievalQA {
 	return NewConversationalRetrievalQA(
 		LoadStuffQA(llm),
@@ -94,8 +93,22 @@ func (c ConversationalRetrievalQA) Call(ctx context.Context, values map[string]a
 	if !ok {
 		return nil, fmt.Errorf("%w: %w", ErrInvalidInputValues, ErrInputValuesWrongType)
 	}
+	chatHistoryStr, ok := values[c.Memory.GetMemoryKey()].(string)
+	if !ok {
+		chatHistory, ok := values[c.Memory.GetMemoryKey()].([]schema.ChatMessage)
+		if !ok {
+			return nil, fmt.Errorf("%w: %w", ErrMissingMemoryKeyValues, ErrMemoryValuesWrongType)
+		}
+
+		bufferStr, err := schema.GetBufferString(chatHistory, "Human", "AI")
+		if err != nil {
+			return nil, err
+		}
+
+		chatHistoryStr = bufferStr
+	}
 
-	question, err := c.getQuestion(ctx, query)
+	question, err := c.getQuestion(ctx, query, chatHistoryStr)
 	if err != nil {
 		return nil, err
 	}
@@ -127,7 +140,7 @@ func (c ConversationalRetrievalQA) Call(ctx context.Context, values map[string]a
 }
 
 func (c ConversationalRetrievalQA) GetMemory() schema.Memory {
-	return &c.Memory
+	return c.Memory
 }
 
 func (c ConversationalRetrievalQA) GetInputKeys() []string {
@@ -143,16 +156,15 @@ func (c ConversationalRetrievalQA) GetOutputKeys() []string {
 	return outputKeys
 }
 
-func (c ConversationalRetrievalQA) getQuestion(ctx context.Context, question string) (string, error) {
-	if len(c.Memory.ChatHistory.Messages()) == 0 {
+func (c ConversationalRetrievalQA) getQuestion(
+	ctx context.Context,
+	question string,
+	chatHistoryStr string,
+) (string, error) {
+	if len(chatHistoryStr) == 0 {
 		return question, nil
 	}
 
-	chatHistoryStr, err := schema.GetBufferString(c.Memory.ChatHistory.Messages(), "Human", "AI")
-	if err != nil {
-		return "", err
-	}
-
 	results, err := Call(
 		ctx,
 		c.CondenseQuestionChain,
diff --git a/chains/errors.go b/chains/errors.go
index 42438c34..d6470f7a 100644
--- a/chains/errors.go
+++ b/chains/errors.go
@@ -11,6 +11,12 @@ var (
 	// ErrInputValuesWrongType is returned if an input value to a chain is of
 	// wrong type.
 	ErrInputValuesWrongType = errors.New("input key is of wrong type")
+	// ErrMissingMemoryKeyValues is returned when some expected input values keys to
+	// a chain is missing.
+	ErrMissingMemoryKeyValues = errors.New("missing memory key in input values")
+	// ErrMemoryValuesWrongType is returned if the memory value to a chain is of
+	// wrong type.
+	ErrMemoryValuesWrongType = errors.New("memory key is of wrong type")
 	// ErrInvalidOutputValues is returned when expected output keys to a chain does
 	// not match the actual keys in the return output values map.
 	ErrInvalidOutputValues = errors.New("missing key in output values")
@@ -21,7 +27,7 @@ var (
 	// ErrMultipleOutputsInRun is returned in the run function if the chain expects
 	// more then one output values.
 	ErrMultipleOutputsInRun = errors.New("run not supported in chain with more then one expected output")
-	// ErrMultipleOutputsInRun is returned in the run function if the chain returns
+	// ErrWrongOutputTypeInRun is returned in the run function if the chain returns
 	// a value that is not a string.
 	ErrWrongOutputTypeInRun = errors.New("run not supported in chain that returns value that is not string")
 
@@ -31,6 +37,6 @@ var (
 	// ErrMultipleOutputsInPredict is returned if a chain has multiple return values
 	// in predict.
 	ErrMultipleOutputsInPredict = errors.New("predict is not supported with a chain that returns multiple values")
-	// ErrChainNotInitialized is returned if a chain is not initialized appropriately.
+	// ErrChainInitialization is returned if a chain is not initialized appropriately.
 	ErrChainInitialization = errors.New("error initializing chain")
 )
diff --git a/memory/buffer.go b/memory/buffer.go
index e4ce2974..54f63f0d 100644
--- a/memory/buffer.go
+++ b/memory/buffer.go
@@ -10,8 +10,8 @@ import (
 // ErrInvalidInputValues is returned when input values given to a memory in save context are invalid.
 var ErrInvalidInputValues = errors.New("invalid input values")
 
-// Buffer is a simple form of memory that remembers previous conversational back and forths directly.
-type Buffer struct {
+// ConversationBuffer is a simple form of memory that remembers previous conversational back and forths directly.
+type ConversationBuffer struct {
 	ChatHistory schema.ChatMessageHistory
 
 	ReturnMessages bool
@@ -22,16 +22,16 @@ type Buffer struct {
 	MemoryKey      string
 }
 
-// Statically assert that Buffer implement the memory interface.
-var _ schema.Memory = &Buffer{}
+// Statically assert that ConversationBuffer implement the memory interface.
+var _ schema.Memory = &ConversationBuffer{}
 
-// NewBuffer is a function for crating a new buffer memory.
-func NewBuffer(options ...BufferOption) *Buffer {
+// NewConversationBuffer is a function for crating a new buffer memory.
+func NewConversationBuffer(options ...ConversationBufferOption) *ConversationBuffer {
 	return applyBufferOptions(options...)
 }
 
 // MemoryVariables gets the input key the buffer memory class will load dynamically.
-func (m *Buffer) MemoryVariables() []string {
+func (m *ConversationBuffer) MemoryVariables() []string {
 	return []string{m.MemoryKey}
 }
 
@@ -39,14 +39,19 @@ func (m *Buffer) MemoryVariables() []string {
 // are returned in a map with the key specified in the MemoryKey field. This key defaults to
 // "history". If ReturnMessages is set to true the output is a slice of schema.ChatMessage. Otherwise
 // the output is a buffer string of the chat messages.
-func (m *Buffer) LoadMemoryVariables(map[string]any) (map[string]any, error) {
+func (m *ConversationBuffer) LoadMemoryVariables(map[string]any) (map[string]any, error) {
+	messages, err := m.ChatHistory.Messages()
+	if err != nil {
+		return nil, err
+	}
+
 	if m.ReturnMessages {
 		return map[string]any{
-			m.MemoryKey: m.ChatHistory.Messages(),
+			m.MemoryKey: messages,
 		}, nil
 	}
 
-	bufferString, err := schema.GetBufferString(m.ChatHistory.Messages(), m.HumanPrefix, m.AIPrefix)
+	bufferString, err := schema.GetBufferString(messages, m.HumanPrefix, m.AIPrefix)
 	if err != nil {
 		return nil, err
 	}
@@ -63,28 +68,35 @@ func (m *Buffer) LoadMemoryVariables(map[string]any) (map[string]any, error) {
 // input key must be a key in the input values and the output key must be a key in the output
 // values. The values in the input and output values used to save a user and ai message must
 // be strings.
-func (m *Buffer) SaveContext(inputValues map[string]any, outputValues map[string]any) error {
+func (m *ConversationBuffer) SaveContext(inputValues map[string]any, outputValues map[string]any) error {
 	userInputValue, err := getInputValue(inputValues, m.InputKey)
 	if err != nil {
 		return err
 	}
-
-	m.ChatHistory.AddUserMessage(userInputValue)
+	err = m.ChatHistory.AddUserMessage(userInputValue)
+	if err != nil {
+		return err
+	}
 
 	aiOutputValue, err := getInputValue(outputValues, m.OutputKey)
 	if err != nil {
 		return err
 	}
-
-	m.ChatHistory.AddAIMessage(aiOutputValue)
+	err = m.ChatHistory.AddAIMessage(aiOutputValue)
+	if err != nil {
+		return err
+	}
 
 	return nil
 }
 
 // Clear sets the chat messages to a new and empty chat message history.
-func (m *Buffer) Clear() error {
-	m.ChatHistory.Clear()
-	return nil
+func (m *ConversationBuffer) Clear() error {
+	return m.ChatHistory.Clear()
+}
+
+func (m *ConversationBuffer) GetMemoryKey() string {
+	return m.MemoryKey
 }
 
 func getInputValue(inputValues map[string]any, inputKey string) (string, error) {
diff --git a/memory/buffer_options.go b/memory/buffer_options.go
index e952e32d..e7564d0c 100644
--- a/memory/buffer_options.go
+++ b/memory/buffer_options.go
@@ -2,61 +2,61 @@ package memory
 
 import "github.com/tmc/langchaingo/schema"
 
-// BufferOption is a function for creating new buffer
+// ConversationBufferOption is a function for creating new buffer
 // with other then the default values.
-type BufferOption func(b *Buffer)
+type ConversationBufferOption func(b *ConversationBuffer)
 
 // WithChatHistory is an option for providing the chat history store.
-func WithChatHistory(chatHistory schema.ChatMessageHistory) BufferOption {
-	return func(b *Buffer) {
+func WithChatHistory(chatHistory schema.ChatMessageHistory) ConversationBufferOption {
+	return func(b *ConversationBuffer) {
 		b.ChatHistory = chatHistory
 	}
 }
 
 // WithReturnMessages is an option for specifying should it return messages.
-func WithReturnMessages(returnMessages bool) BufferOption {
-	return func(b *Buffer) {
+func WithReturnMessages(returnMessages bool) ConversationBufferOption {
+	return func(b *ConversationBuffer) {
 		b.ReturnMessages = returnMessages
 	}
 }
 
 // WithInputKey is an option for specifying the input key.
-func WithInputKey(inputKey string) BufferOption {
-	return func(b *Buffer) {
+func WithInputKey(inputKey string) ConversationBufferOption {
+	return func(b *ConversationBuffer) {
 		b.InputKey = inputKey
 	}
 }
 
 // WithOutputKey is an option for specifying the output key.
-func WithOutputKey(outputKey string) BufferOption {
-	return func(b *Buffer) {
+func WithOutputKey(outputKey string) ConversationBufferOption {
+	return func(b *ConversationBuffer) {
 		b.OutputKey = outputKey
 	}
 }
 
 // WithHumanPrefix is an option for specifying the human prefix.
-func WithHumanPrefix(humanPrefix string) BufferOption {
-	return func(b *Buffer) {
+func WithHumanPrefix(humanPrefix string) ConversationBufferOption {
+	return func(b *ConversationBuffer) {
 		b.HumanPrefix = humanPrefix
 	}
 }
 
 // WithAIPrefix is an option for specifying the AI prefix.
-func WithAIPrefix(aiPrefix string) BufferOption {
-	return func(b *Buffer) {
+func WithAIPrefix(aiPrefix string) ConversationBufferOption {
+	return func(b *ConversationBuffer) {
 		b.AIPrefix = aiPrefix
 	}
 }
 
 // WithMemoryKey is an option for specifying the memory key.
-func WithMemoryKey(memoryKey string) BufferOption {
-	return func(b *Buffer) {
+func WithMemoryKey(memoryKey string) ConversationBufferOption {
+	return func(b *ConversationBuffer) {
 		b.MemoryKey = memoryKey
 	}
 }
 
-func applyBufferOptions(opts ...BufferOption) *Buffer {
-	m := &Buffer{
+func applyBufferOptions(opts ...ConversationBufferOption) *ConversationBuffer {
+	m := &ConversationBuffer{
 		ReturnMessages: false,
 		InputKey:       "",
 		OutputKey:      "",
diff --git a/memory/chat.go b/memory/chat.go
index 1f5798f4..4bd6225a 100644
--- a/memory/chat.go
+++ b/memory/chat.go
@@ -16,28 +16,33 @@ func NewChatMessageHistory(options ...ChatMessageHistoryOption) *ChatMessageHist
 }
 
 // Messages returns all messages stored.
-func (h *ChatMessageHistory) Messages() []schema.ChatMessage {
-	return h.messages
+func (h *ChatMessageHistory) Messages() ([]schema.ChatMessage, error) {
+	return h.messages, nil
 }
 
 // AddAIMessage adds an AIMessage to the chat message history.
-func (h *ChatMessageHistory) AddAIMessage(text string) {
+func (h *ChatMessageHistory) AddAIMessage(text string) error {
 	h.messages = append(h.messages, schema.AIChatMessage{Content: text})
+	return nil
 }
 
 // AddUserMessage adds an user to the chat message history.
-func (h *ChatMessageHistory) AddUserMessage(text string) {
+func (h *ChatMessageHistory) AddUserMessage(text string) error {
 	h.messages = append(h.messages, schema.HumanChatMessage{Content: text})
+	return nil
 }
 
-func (h *ChatMessageHistory) Clear() {
+func (h *ChatMessageHistory) Clear() error {
 	h.messages = make([]schema.ChatMessage, 0)
+	return nil
 }
 
-func (h *ChatMessageHistory) AddMessage(message schema.ChatMessage) {
+func (h *ChatMessageHistory) AddMessage(message schema.ChatMessage) error {
 	h.messages = append(h.messages, message)
+	return nil
 }
 
-func (h *ChatMessageHistory) SetMessages(messages []schema.ChatMessage) {
+func (h *ChatMessageHistory) SetMessages(messages []schema.ChatMessage) error {
 	h.messages = messages
+	return nil
 }
diff --git a/memory/doc.go b/memory/doc.go
index 7027209e..50d86528 100644
--- a/memory/doc.go
+++ b/memory/doc.go
@@ -4,6 +4,6 @@ a variety of implementations for storing and retrieving that data.
 
 The main components of this package are:
 - ChatMessageHistory: a struct that stores chat messages.
-- Buffer: a simple form of memory that remembers previous conversational back and forths directly.
+- ConversationBuffer: a simple form of memory that remembers previous conversational back and forths directly.
 */
 package memory
diff --git a/memory/simple.go b/memory/simple.go
index f4f37c08..244db7e9 100644
--- a/memory/simple.go
+++ b/memory/simple.go
@@ -30,3 +30,7 @@ func (m Simple) SaveContext(map[string]any, map[string]any) error {
 func (m Simple) Clear() error {
 	return nil
 }
+
+func (m Simple) GetMemoryKey() string {
+	return ""
+}
diff --git a/memory/token_buffer.go b/memory/token_buffer.go
index 6244a7c5..8bb0c5ff 100644
--- a/memory/token_buffer.go
+++ b/memory/token_buffer.go
@@ -5,40 +5,44 @@ import (
 	"github.com/tmc/langchaingo/schema"
 )
 
-// TokenBuffer for storing conversation memory.
-type TokenBuffer struct {
-	Buffer
+// ConversationTokenBuffer for storing conversation memory.
+type ConversationTokenBuffer struct {
+	ConversationBuffer
 	LLM           llms.LanguageModel
 	MaxTokenLimit int
 }
 
-// Statically assert that TokenBuffer implement the memory interface.
-var _ schema.Memory = &TokenBuffer{}
+// Statically assert that ConversationTokenBuffer implement the memory interface.
+var _ schema.Memory = &ConversationTokenBuffer{}
 
-// NewTokenBuffer is a function for crating a new token buffer memory.
-func NewTokenBuffer(llm llms.LanguageModel, maxTokenLimit int, options ...BufferOption) *TokenBuffer {
-	tb := &TokenBuffer{
-		LLM:           llm,
-		MaxTokenLimit: maxTokenLimit,
-		Buffer:        *applyBufferOptions(options...),
+// NewConversationTokenBuffer is a function for crating a new token buffer memory.
+func NewConversationTokenBuffer(
+	llm llms.LanguageModel,
+	maxTokenLimit int,
+	options ...ConversationBufferOption,
+) *ConversationTokenBuffer {
+	tb := &ConversationTokenBuffer{
+		LLM:                llm,
+		MaxTokenLimit:      maxTokenLimit,
+		ConversationBuffer: *applyBufferOptions(options...),
 	}
 
 	return tb
 }
 
-// MemoryVariables uses Buffer method for memory variables.
-func (tb *TokenBuffer) MemoryVariables() []string {
-	return tb.Buffer.MemoryVariables()
+// MemoryVariables uses ConversationBuffer method for memory variables.
+func (tb *ConversationTokenBuffer) MemoryVariables() []string {
+	return tb.ConversationBuffer.MemoryVariables()
 }
 
-// LoadMemoryVariables uses Buffer method for loading memory variables.
-func (tb *TokenBuffer) LoadMemoryVariables(inputs map[string]any) (map[string]any, error) {
-	return tb.Buffer.LoadMemoryVariables(inputs)
+// LoadMemoryVariables uses ConversationBuffer method for loading memory variables.
+func (tb *ConversationTokenBuffer) LoadMemoryVariables(inputs map[string]any) (map[string]any, error) {
+	return tb.ConversationBuffer.LoadMemoryVariables(inputs)
 }
 
-// SaveContext uses Buffer method for saving context and prunes memory buffer if needed.
-func (tb *TokenBuffer) SaveContext(inputValues map[string]any, outputValues map[string]any) error {
-	err := tb.Buffer.SaveContext(inputValues, outputValues)
+// SaveContext uses ConversationBuffer method for saving context and prunes memory buffer if needed.
+func (tb *ConversationTokenBuffer) SaveContext(inputValues map[string]any, outputValues map[string]any) error {
+	err := tb.ConversationBuffer.SaveContext(inputValues, outputValues)
 	if err != nil {
 		return err
 	}
@@ -51,11 +55,20 @@ func (tb *TokenBuffer) SaveContext(inputValues map[string]any, outputValues map[
 		// while currBufferLength is greater than MaxTokenLimit we keep removing messages from the memory
 		// from the oldest
 		for currBufferLength > tb.MaxTokenLimit {
-			if len(tb.ChatHistory.Messages()) == 0 {
+			messages, err := tb.ChatHistory.Messages()
+			if err != nil {
+				return err
+			}
+
+			if len(messages) == 0 {
 				break
 			}
 
-			tb.ChatHistory.SetMessages(append(tb.ChatHistory.Messages()[:0], tb.ChatHistory.Messages()[1:]...))
+			err = tb.ChatHistory.SetMessages(append(messages[:0], messages[1:]...))
+			if err != nil {
+				return err
+			}
+
 			currBufferLength, err = tb.getNumTokensFromMessages()
 			if err != nil {
 				return err
@@ -66,13 +79,22 @@ func (tb *TokenBuffer) SaveContext(inputValues map[string]any, outputValues map[
 	return nil
 }
 
-// Clear uses Buffer method for clearing buffer memory.
-func (tb *TokenBuffer) Clear() error {
-	return tb.Buffer.Clear()
+// Clear uses ConversationBuffer method for clearing buffer memory.
+func (tb *ConversationTokenBuffer) Clear() error {
+	return tb.ConversationBuffer.Clear()
 }
 
-func (tb *TokenBuffer) getNumTokensFromMessages() (int, error) {
-	bufferString, err := schema.GetBufferString(tb.ChatHistory.Messages(), tb.Buffer.HumanPrefix, tb.Buffer.AIPrefix)
+func (tb *ConversationTokenBuffer) getNumTokensFromMessages() (int, error) {
+	messages, err := tb.ChatHistory.Messages()
+	if err != nil {
+		return 0, err
+	}
+
+	bufferString, err := schema.GetBufferString(
+		messages,
+		tb.ConversationBuffer.HumanPrefix,
+		tb.ConversationBuffer.AIPrefix,
+	)
 	if err != nil {
 		return 0, err
 	}
diff --git a/schema/chat_message_history.go b/schema/chat_message_history.go
index 2507ee42..2c573e9f 100644
--- a/schema/chat_message_history.go
+++ b/schema/chat_message_history.go
@@ -3,20 +3,20 @@ package schema
 // ChatMessageHistory is the interface for chat history in memory/store.
 type ChatMessageHistory interface {
 	// AddUserMessage Convenience method for adding a human message string to the store.
-	AddUserMessage(message string)
+	AddUserMessage(message string) error
 
 	// AddAIMessage Convenience method for adding an AI message string to the store.
-	AddAIMessage(message string)
+	AddAIMessage(message string) error
 
 	// AddMessage Add a Message object to the store.
-	AddMessage(message ChatMessage)
+	AddMessage(message ChatMessage) error
 
 	// Clear Remove all messages from the store.
-	Clear()
+	Clear() error
 
 	// Messages get all messages from the store
-	Messages() []ChatMessage
+	Messages() ([]ChatMessage, error)
 
 	// SetMessages replaces existing messages in the store
-	SetMessages(messages []ChatMessage)
+	SetMessages(messages []ChatMessage) error
 }
diff --git a/schema/memory.go b/schema/memory.go
index a6c75e2b..c5622163 100644
--- a/schema/memory.go
+++ b/schema/memory.go
@@ -2,6 +2,8 @@ package schema
 
 // Memory is the interface for memory in chains.
 type Memory interface {
+	// GetMemoryKey getter for memory key.
+	GetMemoryKey() string
 	// MemoryVariables Input keys this memory class will load dynamically.
 	MemoryVariables() []string
 	// LoadMemoryVariables Return key-value pairs given the text input to the chain.


Test patch:
diff --git a/agents/conversational_test.go b/agents/conversational_test.go
index a895d229..72ec8de0 100644
--- a/agents/conversational_test.go
+++ b/agents/conversational_test.go
@@ -26,7 +26,7 @@ func TestConversationalWithMemory(t *testing.T) {
 		llm,
 		[]tools.Tool{tools.Calculator{}},
 		ConversationalReactDescription,
-		WithMemory(memory.NewBuffer()),
+		WithMemory(memory.NewConversationBuffer()),
 	)
 	require.NoError(t, err)
 
diff --git a/chains/conversation_test.go b/chains/conversation_test.go
index cd36960e..c5f74d74 100644
--- a/chains/conversation_test.go
+++ b/chains/conversation_test.go
@@ -20,7 +20,7 @@ func TestConversation(t *testing.T) {
 	llm, err := openai.New()
 	require.NoError(t, err)
 
-	c := NewConversation(llm, memory.NewBuffer())
+	c := NewConversation(llm, memory.NewConversationBuffer())
 	_, err = Run(context.Background(), c, "Hi! I'm Jim")
 	require.NoError(t, err)
 
@@ -39,7 +39,7 @@ func TestConversationMemoryPrune(t *testing.T) {
 	llm, err := openai.New()
 	require.NoError(t, err)
 
-	c := NewConversation(llm, memory.NewTokenBuffer(llm, 50))
+	c := NewConversation(llm, memory.NewConversationTokenBuffer(llm, 50))
 	_, err = Run(context.Background(), c, "Hi! I'm Jim")
 	require.NoError(t, err)
 
diff --git a/chains/conversational_retrieval_qa_test.go b/chains/conversational_retrieval_qa_test.go
index d262a8c1..7e8fa1f4 100644
--- a/chains/conversational_retrieval_qa_test.go
+++ b/chains/conversational_retrieval_qa_test.go
@@ -73,14 +73,49 @@ func TestConversationalRetrievalQA(t *testing.T) {
 	combinedQuestionGeneratorChain := LoadCondenseQuestionGenerator(llm)
 	r := testConversationalRetriever{}
 
-	chain := NewConversationalRetrievalQA(combinedStuffQAChain, combinedQuestionGeneratorChain, r, *memory.NewBuffer())
+	chain := NewConversationalRetrievalQA(
+		combinedStuffQAChain,
+		combinedQuestionGeneratorChain,
+		r,
+		memory.NewConversationBuffer(),
+	)
 	result, err := Run(ctx, chain, "What did the president say about Ketanji Brown Jackson")
 	require.NoError(t, err)
 	require.True(t, strings.Contains(result, "Ketanji Brown Jackson"), "expected Ketanji Brown Jackson in result")
 
 	result, err = Run(ctx, chain, "Did he mention who she succeeded")
 	require.NoError(t, err)
-	require.True(t, strings.Contains(result, " Justice Stephen Breyer"), "expected  Justice Stephen Breyer in result")
+	require.True(t, strings.Contains(result, "Justice Stephen Breyer"), "expected  Justice Stephen Breyer in result")
+}
+
+func TestConversationalRetrievalQAWithReturnMessages(t *testing.T) {
+	t.Parallel()
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+
+	ctx := context.Background()
+
+	llm, err := openai.New()
+	require.NoError(t, err)
+
+	combinedStuffQAChain := LoadStuffQA(llm)
+	combinedQuestionGeneratorChain := LoadCondenseQuestionGenerator(llm)
+	r := testConversationalRetriever{}
+
+	chain := NewConversationalRetrievalQA(
+		combinedStuffQAChain,
+		combinedQuestionGeneratorChain,
+		r,
+		memory.NewConversationBuffer(memory.WithReturnMessages(true)),
+	)
+	result, err := Run(ctx, chain, "What did the president say about Ketanji Brown Jackson")
+	require.NoError(t, err)
+	require.True(t, strings.Contains(result, "Ketanji Brown Jackson"), "expected Ketanji Brown Jackson in result")
+
+	result, err = Run(ctx, chain, "Did he mention who she succeeded")
+	require.NoError(t, err)
+	require.True(t, strings.Contains(result, "Justice Stephen Breyer"), "expected  Justice Stephen Breyer in result")
 }
 
 func TestConversationalRetrievalQAFromLLM(t *testing.T) {
@@ -95,12 +130,38 @@ func TestConversationalRetrievalQAFromLLM(t *testing.T) {
 	llm, err := openai.New()
 	require.NoError(t, err)
 
-	chain := NewConversationalRetrievalQAFromLLM(llm, r, *memory.NewBuffer())
+	chain := NewConversationalRetrievalQAFromLLM(llm, r, memory.NewConversationBuffer())
+	result, err := Run(context.Background(), chain, "What did the president say about Ketanji Brown Jackson")
+	require.NoError(t, err)
+	require.True(t, strings.Contains(result, "Ketanji Brown Jackson"), "expected Ketanji Brown Jackson in result")
+
+	result, err = Run(ctx, chain, "Did he mention who she succeeded")
+	require.NoError(t, err)
+	require.True(t, strings.Contains(result, "Justice Stephen Breyer"), "expected  Justice Stephen Breyer in result")
+}
+
+func TestConversationalRetrievalQAFromLLMWithConversationTokenBuffer(t *testing.T) {
+	t.Parallel()
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+
+	ctx := context.Background()
+
+	r := testConversationalRetriever{}
+	llm, err := openai.New()
+	require.NoError(t, err)
+
+	chain := NewConversationalRetrievalQAFromLLM(
+		llm,
+		r,
+		memory.NewConversationTokenBuffer(llm, 2000),
+	)
 	result, err := Run(context.Background(), chain, "What did the president say about Ketanji Brown Jackson")
 	require.NoError(t, err)
 	require.True(t, strings.Contains(result, "Ketanji Brown Jackson"), "expected Ketanji Brown Jackson in result")
 
 	result, err = Run(ctx, chain, "Did he mention who she succeeded")
 	require.NoError(t, err)
-	require.True(t, strings.Contains(result, " Justice Stephen Breyer"), "expected  Justice Stephen Breyer in result")
+	require.True(t, strings.Contains(result, "Justice Stephen Breyer"), "expected  Justice Stephen Breyer in result")
 }
diff --git a/chains/sequential_test.go b/chains/sequential_test.go
index 27f65484..5dbb7b28 100644
--- a/chains/sequential_test.go
+++ b/chains/sequential_test.go
@@ -171,8 +171,10 @@ func TestSequentialChainErrors(t *testing.T) {
 			chains: []Chain{
 				&testLLMChain{inputKeys: []string{"input1"}, outputKeys: []string{"output"}},
 			},
-			initErr:      ErrChainInitialization,
-			seqChainOpts: []SequentialChainOption{WithSeqChainMemory(memory.NewBuffer(memory.WithMemoryKey("input1")))},
+			initErr: ErrChainInitialization,
+			seqChainOpts: []SequentialChainOption{WithSeqChainMemory(
+				memory.NewConversationBuffer(memory.WithMemoryKey("input1")),
+			)},
 		},
 	}
 
diff --git a/memory/buffer_test.go b/memory/buffer_test.go
index 3a6c2b41..97562269 100644
--- a/memory/buffer_test.go
+++ b/memory/buffer_test.go
@@ -11,7 +11,7 @@ import (
 func TestBufferMemory(t *testing.T) {
 	t.Parallel()
 
-	m := NewBuffer()
+	m := NewConversationBuffer()
 	result1, err := m.LoadMemoryVariables(map[string]any{})
 	require.NoError(t, err)
 	expected1 := map[string]any{"history": ""}
@@ -30,7 +30,7 @@ func TestBufferMemory(t *testing.T) {
 func TestBufferMemoryReturnMessage(t *testing.T) {
 	t.Parallel()
 
-	m := NewBuffer()
+	m := NewConversationBuffer()
 	m.ReturnMessages = true
 	expected1 := map[string]any{"history": []schema.ChatMessage{}}
 	result1, err := m.LoadMemoryVariables(map[string]any{})
@@ -50,14 +50,16 @@ func TestBufferMemoryReturnMessage(t *testing.T) {
 		}),
 	)
 
-	expected2 := map[string]any{"history": expectedChatHistory.Messages()}
+	messages, err := expectedChatHistory.Messages()
+	assert.NoError(t, err)
+	expected2 := map[string]any{"history": messages}
 	assert.Equal(t, expected2, result2)
 }
 
 func TestBufferMemoryWithPreLoadedHistory(t *testing.T) {
 	t.Parallel()
 
-	m := NewBuffer(WithChatHistory(NewChatMessageHistory(
+	m := NewConversationBuffer(WithChatHistory(NewChatMessageHistory(
 		WithPreviousMessages([]schema.ChatMessage{
 			schema.HumanChatMessage{Content: "bar"},
 			schema.AIChatMessage{Content: "foo"},
@@ -74,33 +76,38 @@ type testChatMessageHistory struct{}
 
 var _ schema.ChatMessageHistory = testChatMessageHistory{}
 
-func (t testChatMessageHistory) AddUserMessage(_ string) {
+func (t testChatMessageHistory) AddUserMessage(_ string) error {
+	return nil
 }
 
-func (t testChatMessageHistory) AddAIMessage(_ string) {
+func (t testChatMessageHistory) AddAIMessage(_ string) error {
+	return nil
 }
 
-func (t testChatMessageHistory) AddMessage(_ schema.ChatMessage) {
+func (t testChatMessageHistory) AddMessage(_ schema.ChatMessage) error {
+	return nil
 }
 
-func (t testChatMessageHistory) Clear() {
+func (t testChatMessageHistory) Clear() error {
+	return nil
 }
 
-func (t testChatMessageHistory) SetMessages(_ []schema.ChatMessage) {
+func (t testChatMessageHistory) SetMessages(_ []schema.ChatMessage) error {
+	return nil
 }
 
-func (t testChatMessageHistory) Messages() []schema.ChatMessage {
+func (t testChatMessageHistory) Messages() ([]schema.ChatMessage, error) {
 	return []schema.ChatMessage{
 		schema.HumanChatMessage{Content: "user message test"},
 		schema.AIChatMessage{Content: "ai message test"},
-	}
+	}, nil
 }
 
 func TestBufferMemoryWithChatHistoryOption(t *testing.T) {
 	t.Parallel()
 
 	chatMessageHistory := testChatMessageHistory{}
-	m := NewBuffer(WithChatHistory(chatMessageHistory))
+	m := NewConversationBuffer(WithChatHistory(chatMessageHistory))
 
 	result, err := m.LoadMemoryVariables(map[string]any{})
 	require.NoError(t, err)
diff --git a/memory/chat_test.go b/memory/chat_test.go
index 59c7528e..c6e7122d 100644
--- a/memory/chat_test.go
+++ b/memory/chat_test.go
@@ -11,12 +11,18 @@ func TestChatMessageHistory(t *testing.T) {
 	t.Parallel()
 
 	h := NewChatMessageHistory()
-	h.AddAIMessage("foo")
-	h.AddUserMessage("bar")
+	err := h.AddAIMessage("foo")
+	assert.NoError(t, err)
+	err = h.AddUserMessage("bar")
+	assert.NoError(t, err)
+
+	messages, err := h.Messages()
+	assert.NoError(t, err)
+
 	assert.Equal(t, []schema.ChatMessage{
 		schema.AIChatMessage{Content: "foo"},
 		schema.HumanChatMessage{Content: "bar"},
-	}, h.Messages())
+	}, messages)
 
 	h = NewChatMessageHistory(
 		WithPreviousMessages([]schema.ChatMessage{
@@ -24,10 +30,15 @@ func TestChatMessageHistory(t *testing.T) {
 			schema.SystemChatMessage{Content: "bar"},
 		}),
 	)
-	h.AddUserMessage("zoo")
+	err = h.AddUserMessage("zoo")
+	assert.NoError(t, err)
+
+	messages, err = h.Messages()
+	assert.NoError(t, err)
+
 	assert.Equal(t, []schema.ChatMessage{
 		schema.AIChatMessage{Content: "foo"},
 		schema.SystemChatMessage{Content: "bar"},
 		schema.HumanChatMessage{Content: "zoo"},
-	}, h.Messages())
+	}, messages)
 }
diff --git a/memory/token_buffer_test.go b/memory/token_buffer_test.go
index 4ea4b2f2..2a93f249 100644
--- a/memory/token_buffer_test.go
+++ b/memory/token_buffer_test.go
@@ -19,7 +19,7 @@ func TestTokenBufferMemory(t *testing.T) {
 
 	llm, err := openai.New()
 	require.NoError(t, err)
-	m := NewTokenBuffer(llm, 2000)
+	m := NewConversationTokenBuffer(llm, 2000)
 
 	result1, err := m.LoadMemoryVariables(map[string]any{})
 	require.NoError(t, err)
@@ -45,7 +45,7 @@ func TestTokenBufferMemoryReturnMessage(t *testing.T) {
 
 	llm, err := openai.New()
 	require.NoError(t, err)
-	m := NewTokenBuffer(llm, 2000, WithReturnMessages(true))
+	m := NewConversationTokenBuffer(llm, 2000, WithReturnMessages(true))
 
 	expected1 := map[string]any{"history": []schema.ChatMessage{}}
 	result1, err := m.LoadMemoryVariables(map[string]any{})
@@ -65,7 +65,9 @@ func TestTokenBufferMemoryReturnMessage(t *testing.T) {
 		}),
 	)
 
-	expected2 := map[string]any{"history": expectedChatHistory.Messages()}
+	messages, err := expectedChatHistory.Messages()
+	require.NoError(t, err)
+	expected2 := map[string]any{"history": messages}
 	assert.Equal(t, expected2, result2)
 }
 
@@ -79,7 +81,7 @@ func TestTokenBufferMemoryWithPreLoadedHistory(t *testing.T) {
 	llm, err := openai.New()
 	require.NoError(t, err)
 
-	m := NewTokenBuffer(llm, 2000, WithChatHistory(NewChatMessageHistory(
+	m := NewConversationTokenBuffer(llm, 2000, WithChatHistory(NewChatMessageHistory(
 		WithPreviousMessages([]schema.ChatMessage{
 			schema.HumanChatMessage{Content: "bar"},
 			schema.AIChatMessage{Content: "foo"},