Repository: tmc/langchaingo
Pull request #209 (2023-07-21 16:15:38)

feat: Memory Integrations (chat history)
### 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
- adding the ability to port in any memory integration for the buffer
- by this logic, we can pass in any store that we would like to have to keep track of the chat history
![Screenshot 2023-07-21 at 6 00 32 PM](https://github.com/tmc/langchaingo/assets/3450224/c4d4204f-205d-49d1-ab29-055aa084965c)



Patch:
diff --git a/memory/buffer.go b/memory/buffer.go
index 1dfc4367..8a60a07c 100644
--- a/memory/buffer.go
+++ b/memory/buffer.go
@@ -12,7 +12,7 @@ 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 {
-	ChatHistory *ChatMessageHistory
+	chatHistory schema.ChatMessageHistory
 
 	ReturnMessages bool
 	InputKey       string
@@ -26,18 +26,8 @@ type Buffer struct {
 var _ schema.Memory = &Buffer{}
 
 // NewBuffer is a function for crating a new buffer memory.
-func NewBuffer() *Buffer {
-	m := Buffer{
-		ChatHistory:    NewChatMessageHistory(),
-		ReturnMessages: false,
-		InputKey:       "",
-		OutputKey:      "",
-		HumanPrefix:    "Human",
-		AIPrefix:       "AI",
-		MemoryKey:      "history",
-	}
-
-	return &m
+func NewBuffer(options ...BufferOption) *Buffer {
+	return applyBufferOptions(options...)
 }
 
 // MemoryVariables gets the input key the buffer memory class will load dynamically.
@@ -52,11 +42,11 @@ func (m *Buffer) MemoryVariables() []string {
 func (m *Buffer) LoadMemoryVariables(map[string]any) (map[string]any, error) {
 	if m.ReturnMessages {
 		return map[string]any{
-			m.MemoryKey: m.ChatHistory.messages,
+			m.MemoryKey: m.chatHistory.Messages(),
 		}, nil
 	}
 
-	bufferString, err := schema.GetBufferString(m.ChatHistory.messages, m.HumanPrefix, m.AIPrefix)
+	bufferString, err := schema.GetBufferString(m.chatHistory.Messages(), m.HumanPrefix, m.AIPrefix)
 	if err != nil {
 		return nil, err
 	}
@@ -79,21 +69,21 @@ func (m *Buffer) SaveContext(inputValues map[string]any, outputValues map[string
 		return err
 	}
 
-	m.ChatHistory.AddUserMessage(userInputValue)
+	m.chatHistory.AddUserMessage(userInputValue)
 
 	aiOutputValue, err := getInputValue(outputValues, m.OutputKey)
 	if err != nil {
 		return err
 	}
 
-	m.ChatHistory.AddAIMessage(aiOutputValue)
+	m.chatHistory.AddAIMessage(aiOutputValue)
 
 	return nil
 }
 
 // Clear sets the chat messages to a new and empty chat message history.
 func (m *Buffer) Clear() error {
-	m.ChatHistory.Clear()
+	m.chatHistory.Clear()
 	return nil
 }
 
diff --git a/memory/buffer_options.go b/memory/buffer_options.go
new file mode 100644
index 00000000..95a0317f
--- /dev/null
+++ b/memory/buffer_options.go
@@ -0,0 +1,77 @@
+package memory
+
+import "github.com/tmc/langchaingo/schema"
+
+// BufferOption is a function for creating new buffer
+// with other then the default values.
+type BufferOption func(b *Buffer)
+
+// WithChatHistory is an option for providing the chat history store.
+func WithChatHistory(chatHistory schema.ChatMessageHistory) BufferOption {
+	return func(b *Buffer) {
+		b.chatHistory = chatHistory
+	}
+}
+
+// WithReturnMessages is an option for specifying should it return messages.
+func WithReturnMessages(returnMessages bool) BufferOption {
+	return func(b *Buffer) {
+		b.ReturnMessages = returnMessages
+	}
+}
+
+// WithInputKey is an option for specifying the input key.
+func WithInputKey(inputKey string) BufferOption {
+	return func(b *Buffer) {
+		b.InputKey = inputKey
+	}
+}
+
+// WithOutputKey is an option for specifying the output key.
+func WithOutputKey(outputKey string) BufferOption {
+	return func(b *Buffer) {
+		b.OutputKey = outputKey
+	}
+}
+
+// WithHumanPrefix is an option for specifying the human prefix.
+func WithHumanPrefix(humanPrefix string) BufferOption {
+	return func(b *Buffer) {
+		b.HumanPrefix = humanPrefix
+	}
+}
+
+// WithAIPrefix is an option for specifying the AI prefix.
+func WithAIPrefix(aiPrefix string) BufferOption {
+	return func(b *Buffer) {
+		b.AIPrefix = aiPrefix
+	}
+}
+
+// WithMemoryKey is an option for specifying the memory key.
+func WithMemoryKey(memoryKey string) BufferOption {
+	return func(b *Buffer) {
+		b.MemoryKey = memoryKey
+	}
+}
+
+func applyBufferOptions(opts ...BufferOption) *Buffer {
+	m := &Buffer{
+		ReturnMessages: false,
+		InputKey:       "",
+		OutputKey:      "",
+		HumanPrefix:    "Human",
+		AIPrefix:       "AI",
+		MemoryKey:      "history",
+	}
+
+	for _, opt := range opts {
+		opt(m)
+	}
+
+	if m.chatHistory == nil {
+		m.chatHistory = NewChatMessageHistory()
+	}
+
+	return m
+}
diff --git a/memory/chat.go b/memory/chat.go
index 86edc11f..b116e89b 100644
--- a/memory/chat.go
+++ b/memory/chat.go
@@ -7,17 +7,12 @@ type ChatMessageHistory struct {
 	messages []schema.ChatMessage
 }
 
-// NewChatMessageHistory creates a new ChatMessageHistory using chat message options.
-func NewChatMessageHistory(options ...NewChatMessageOption) *ChatMessageHistory {
-	h := &ChatMessageHistory{
-		messages: make([]schema.ChatMessage, 0),
-	}
-
-	for _, option := range options {
-		option(h)
-	}
+// Statically assert that ChatMessageHistory implement the chat message history interface.
+var _ schema.ChatMessageHistory = &ChatMessageHistory{}
 
-	return h
+// NewChatMessageHistory creates a new ChatMessageHistory using chat message options.
+func NewChatMessageHistory(options ...ChatMessageHistoryOption) *ChatMessageHistory {
+	return applyChatOptions(options...)
 }
 
 // Messages returns all messages stored.
@@ -39,14 +34,6 @@ func (h *ChatMessageHistory) Clear() {
 	h.messages = make([]schema.ChatMessage, 0)
 }
 
-// NewChatMessageOption is a function for creating new chat message history
-// with other then the default values.
-type NewChatMessageOption func(m *ChatMessageHistory)
-
-// WithPreviousMessages is an option for NewChatMessageHistory for adding
-// previous messages to the history.
-func WithPreviousMessages(previousMessages []schema.ChatMessage) NewChatMessageOption {
-	return func(m *ChatMessageHistory) {
-		m.messages = append(m.messages, previousMessages...)
-	}
+func (h *ChatMessageHistory) AddMessage(message schema.ChatMessage) {
+	h.messages = append(h.messages, message)
 }
diff --git a/memory/chat_options.go b/memory/chat_options.go
new file mode 100644
index 00000000..a374531f
--- /dev/null
+++ b/memory/chat_options.go
@@ -0,0 +1,27 @@
+package memory
+
+import "github.com/tmc/langchaingo/schema"
+
+// ChatMessageHistoryOption is a function for creating new chat message history
+// with other then the default values.
+type ChatMessageHistoryOption func(m *ChatMessageHistory)
+
+// WithPreviousMessages is an option for NewChatMessageHistory for adding
+// previous messages to the history.
+func WithPreviousMessages(previousMessages []schema.ChatMessage) ChatMessageHistoryOption {
+	return func(m *ChatMessageHistory) {
+		m.messages = append(m.messages, previousMessages...)
+	}
+}
+
+func applyChatOptions(options ...ChatMessageHistoryOption) *ChatMessageHistory {
+	h := &ChatMessageHistory{
+		messages: make([]schema.ChatMessage, 0),
+	}
+
+	for _, option := range options {
+		option(h)
+	}
+
+	return h
+}
diff --git a/schema/chat_message_history.go b/schema/chat_message_history.go
new file mode 100644
index 00000000..1fdde7a0
--- /dev/null
+++ b/schema/chat_message_history.go
@@ -0,0 +1,19 @@
+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)
+
+	// AddAIMessage Convenience method for adding an AI message string to the store.
+	AddAIMessage(message string)
+
+	// AddMessage Add a Message object to the store.
+	AddMessage(message ChatMessage)
+
+	// Clear Remove all messages from the store.
+	Clear()
+
+	// Messages get all messages from the store
+	Messages() []ChatMessage
+}


Test patch:
diff --git a/memory/buffer_test.go b/memory/buffer_test.go
index 358a74af..23d3c701 100644
--- a/memory/buffer_test.go
+++ b/memory/buffer_test.go
@@ -57,16 +57,50 @@ func TestBufferMemoryReturnMessage(t *testing.T) {
 func TestBufferMemoryWithPreLoadedHistory(t *testing.T) {
 	t.Parallel()
 
-	m := NewBuffer()
-	m.ChatHistory = NewChatMessageHistory(
+	m := NewBuffer(WithChatHistory(NewChatMessageHistory(
 		WithPreviousMessages([]schema.ChatMessage{
 			schema.HumanChatMessage{Text: "bar"},
 			schema.AIChatMessage{Text: "foo"},
 		}),
-	)
+	)))
 
 	result, err := m.LoadMemoryVariables(map[string]any{})
 	require.NoError(t, err)
 	expected := map[string]any{"history": "Human: bar\nAI: foo"}
 	assert.Equal(t, expected, result)
 }
+
+type testChatMessageHistory struct{}
+
+var _ schema.ChatMessageHistory = testChatMessageHistory{}
+
+func (t testChatMessageHistory) AddUserMessage(_ string) {
+}
+
+func (t testChatMessageHistory) AddAIMessage(_ string) {
+}
+
+func (t testChatMessageHistory) AddMessage(_ schema.ChatMessage) {
+}
+
+func (t testChatMessageHistory) Clear() {
+}
+
+func (t testChatMessageHistory) Messages() []schema.ChatMessage {
+	return []schema.ChatMessage{
+		schema.HumanChatMessage{Text: "user message test"},
+		schema.AIChatMessage{Text: "ai message test"},
+	}
+}
+
+func TestBufferMemoryWithChatHistoryOption(t *testing.T) {
+	t.Parallel()
+
+	chatMessageHistory := testChatMessageHistory{}
+	m := NewBuffer(WithChatHistory(chatMessageHistory))
+
+	result, err := m.LoadMemoryVariables(map[string]any{})
+	require.NoError(t, err)
+	expected := map[string]any{"history": "Human: user message test\nAI: ai message test"}
+	assert.Equal(t, expected, result)
+}