From 6710273eac8bbc35cbcf5e75308fb7cd5707bb26 Mon Sep 17 00:00:00 2001 From: zivkovicn Date: Fri, 21 Jul 2023 17:58:17 +0200 Subject: [PATCH] feat-adding-memory-buffer-options | updating memory buffer options --- memory/buffer.go | 20 +++------ memory/buffer_options.go | 77 ++++++++++++++++++++++++++++++++++ memory/buffer_test.go | 36 ++++++++++++++++ memory/chat.go | 27 ++++-------- memory/chat_options.go | 27 ++++++++++++ schema/chat_message_history.go | 19 +++++++++ 6 files changed, 171 insertions(+), 35 deletions(-) create mode 100644 memory/buffer_options.go create mode 100644 memory/chat_options.go create mode 100644 schema/chat_message_history.go diff --git a/memory/buffer.go b/memory/buffer.go index 1dfc4367..e4ce2974 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 } diff --git a/memory/buffer_options.go b/memory/buffer_options.go new file mode 100644 index 00000000..e952e32d --- /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/buffer_test.go b/memory/buffer_test.go index 358a74af..45c9db9f 100644 --- a/memory/buffer_test.go +++ b/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) +} 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 } +// 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) } 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 +}