feat-adding-memory-buffer-options | updating memory buffer options

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
 	}

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
+}

memory/buffer_test.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
 	"github.com/tmc/langchaingo/schema"
 )
 
@@ -70,3 +71,38 @@ func TestBufferMemoryWithPreLoadedHistory(t *testing.T) {
 	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(message string) {
+}
+
+func (t testChatMessageHistory) AddAIMessage(message string) {
+}
+
+func (t testChatMessageHistory) AddMessage(message 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)
+}

memory/chat.go

@@ -7,17 +7,12 @@ type ChatMessageHistory struct {
 	messages []schema.ChatMessage
 }
 
+// Statically assert that ChatMessageHistory implement the chat message history interface.
+var _ schema.ChatMessageHistory = &ChatMessageHistory{}
+
 // 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)
-	}
-
-	return h
+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)
 }

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
+}

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
+}