Repository: tmc/langchaingo
Pull request #60 (2023-04-25 11:13:23)

memory: add version outside of exp
 - Add buffer memory
https://github.com/hwchase17/langchainjs/blob/main/langchain/src/memory/buffer_memory.ts
https://github.com/hwchase17/langchain/blob/master/langchain/memory/chat_memory.py
 - Add chat message history
https://github.com/hwchase17/langchainjs/blob/main/langchain/src/stores/message/in_memory.ts
https://github.com/hwchase17/langchain/blob/master/langchain/memory/chat_message_histories/in_memory.py

### 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 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 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.
- [x]  I am an existing contributor.



Patch:
diff --git a/memory/buffer_memory.go b/memory/buffer_memory.go
new file mode 100644
index 00000000..1cb0ca3a
--- /dev/null
+++ b/memory/buffer_memory.go
@@ -0,0 +1,144 @@
+package memory
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/tmc/langchaingo/schema"
+)
+
+// 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 {
+	ChatHistory *ChatMessageHistory
+
+	ReturnMessages bool
+	InputKey       string
+	OutputKey      string
+	HumanPrefix    string
+	AIPrefix       string
+	MemoryKey      string
+}
+
+// Statically assert that Buffer implement the memory interface.
+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
+}
+
+// MemoryVariables gets the input key the buffer memory class will load dynamically.
+func (m *Buffer) MemoryVariables() []string {
+	return []string{m.MemoryKey}
+}
+
+// Returns the previous chat messages stored in memory. Previous chat messages 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) {
+	if m.ReturnMessages {
+		return map[string]any{
+			m.MemoryKey: m.ChatHistory.messages,
+		}, nil
+	}
+
+	bufferString, err := schema.GetBufferString(m.ChatHistory.messages, m.HumanPrefix, m.AIPrefix)
+	if err != nil {
+		return nil, err
+	}
+
+	return map[string]any{
+		m.MemoryKey: bufferString,
+	}, nil
+}
+
+// SaveContext uses the input values to the llm to save a user message, and the output values
+// of the llm to save a ai message. If the input or output key is not set, the input values or
+// output values must contain only one key such that the function can know what string to
+// add as a user and AI message. On the other hand, if the output key or input key is set, the
+// 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 {
+	userInputValue, err := getInputValue(inputValues, m.InputKey)
+	if err != nil {
+		return err
+	}
+
+	m.ChatHistory.AddUserMessage(userInputValue)
+
+	aiOutputValue, err := getInputValue(outputValues, m.OutputKey)
+	if err != nil {
+		return err
+	}
+
+	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()
+	return nil
+}
+
+func getInputValue(inputValues map[string]any, inputKey string) (string, error) {
+	// If the input key is set, return the value in the inputValues with the input key.
+	if inputKey != "" {
+		inputValue, ok := inputValues[inputKey]
+		if !ok {
+			return "", fmt.Errorf(
+				"%w: %v do not contain inputKey %s",
+				ErrInvalidInputValues,
+				inputValues,
+				inputKey,
+			)
+		}
+
+		return getInputValueReturnToString(inputValue)
+	}
+
+	// Otherwise error if length of map isn't one, or return the only entry in the map.
+	if len(inputValues) > 1 {
+		return "", fmt.Errorf(
+			"%w: multiple keys and no input key set",
+			ErrInvalidInputValues,
+		)
+	}
+
+	for _, inputValue := range inputValues {
+		return getInputValueReturnToString(inputValue)
+	}
+
+	return "", fmt.Errorf("%w: 0 keys", ErrInvalidInputValues)
+}
+
+func getInputValueReturnToString(
+	inputValue interface{},
+) (string, error) {
+	switch value := inputValue.(type) {
+	case string:
+		return value, nil
+	default:
+		return "", fmt.Errorf(
+			"%w: input value %v not string",
+			ErrInvalidInputValues,
+			inputValue,
+		)
+	}
+}
diff --git a/memory/chat_memory.go b/memory/chat_memory.go
new file mode 100644
index 00000000..86edc11f
--- /dev/null
+++ b/memory/chat_memory.go
@@ -0,0 +1,52 @@
+package memory
+
+import "github.com/tmc/langchaingo/schema"
+
+// ChatMessageHistory is a struct that stores chat messages.
+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)
+	}
+
+	return h
+}
+
+// Messages returns all messages stored.
+func (h *ChatMessageHistory) Messages() []schema.ChatMessage {
+	return h.messages
+}
+
+// AddAIMessage adds an AIMessage to the chat message history.
+func (h *ChatMessageHistory) AddAIMessage(text string) {
+	h.messages = append(h.messages, schema.AIChatMessage{Text: text})
+}
+
+// AddUserMessage adds an user to the chat message history.
+func (h *ChatMessageHistory) AddUserMessage(text string) {
+	h.messages = append(h.messages, schema.HumanChatMessage{Text: text})
+}
+
+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...)
+	}
+}
diff --git a/memory/doc.go b/memory/doc.go
new file mode 100644
index 00000000..cfdf615c
--- /dev/null
+++ b/memory/doc.go
@@ -0,0 +1,3 @@
+// Package memory contains implementations of the memory interface.
+// Memory is used for storing and retrieving data in the process of a conversation.
+package memory
diff --git a/memory/empty_memory.go b/memory/empty_memory.go
new file mode 100644
index 00000000..372cfd29
--- /dev/null
+++ b/memory/empty_memory.go
@@ -0,0 +1,32 @@
+package memory
+
+import (
+	"github.com/tmc/langchaingo/schema"
+)
+
+// Simple is a class that implement the memory interface, but does nothing.
+// The class is used as default in multiple chains.
+type Simple struct{}
+
+func NewEmpty() Simple {
+	return Simple{}
+}
+
+// Statically assert that Simple implement the memory interface.
+var _ schema.Memory = Simple{}
+
+func (m Simple) MemoryVariables() []string {
+	return nil
+}
+
+func (m Simple) LoadMemoryVariables(map[string]any) (map[string]any, error) {
+	return make(map[string]any, 0), nil
+}
+
+func (m Simple) SaveContext(map[string]any, map[string]any) error {
+	return nil
+}
+
+func (m Simple) Clear() error {
+	return nil
+}


Test patch:
diff --git a/memory/buffer_memory_test.go b/memory/buffer_memory_test.go
new file mode 100644
index 00000000..358a74af
--- /dev/null
+++ b/memory/buffer_memory_test.go
@@ -0,0 +1,72 @@
+package memory
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/tmc/langchaingo/schema"
+)
+
+func TestBufferMemory(t *testing.T) {
+	t.Parallel()
+
+	m := NewBuffer()
+	result1, err := m.LoadMemoryVariables(map[string]any{})
+	require.NoError(t, err)
+	expected1 := map[string]any{"history": ""}
+	assert.Equal(t, expected1, result1)
+
+	err = m.SaveContext(map[string]any{"foo": "bar"}, map[string]any{"bar": "foo"})
+	require.NoError(t, err)
+
+	result2, err := m.LoadMemoryVariables(map[string]any{})
+	require.NoError(t, err)
+
+	expected2 := map[string]any{"history": "Human: bar\nAI: foo"}
+	assert.Equal(t, expected2, result2)
+}
+
+func TestBufferMemoryReturnMessage(t *testing.T) {
+	t.Parallel()
+
+	m := NewBuffer()
+	m.ReturnMessages = true
+	expected1 := map[string]any{"history": []schema.ChatMessage{}}
+	result1, err := m.LoadMemoryVariables(map[string]any{})
+	require.NoError(t, err)
+	assert.Equal(t, expected1, result1)
+
+	err = m.SaveContext(map[string]any{"foo": "bar"}, map[string]any{"bar": "foo"})
+	require.NoError(t, err)
+
+	result2, err := m.LoadMemoryVariables(map[string]any{})
+	require.NoError(t, err)
+
+	expectedChatHistory := NewChatMessageHistory(
+		WithPreviousMessages([]schema.ChatMessage{
+			schema.HumanChatMessage{Text: "bar"},
+			schema.AIChatMessage{Text: "foo"},
+		}),
+	)
+
+	expected2 := map[string]any{"history": expectedChatHistory.Messages()}
+	assert.Equal(t, expected2, result2)
+}
+
+func TestBufferMemoryWithPreLoadedHistory(t *testing.T) {
+	t.Parallel()
+
+	m := NewBuffer()
+	m.ChatHistory = 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)
+}
diff --git a/memory/chat_memory_test.go b/memory/chat_memory_test.go
new file mode 100644
index 00000000..bc503a4c
--- /dev/null
+++ b/memory/chat_memory_test.go
@@ -0,0 +1,33 @@
+package memory
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tmc/langchaingo/schema"
+)
+
+func TestChatMessageHistory(t *testing.T) {
+	t.Parallel()
+
+	h := NewChatMessageHistory()
+	h.AddAIMessage("foo")
+	h.AddUserMessage("bar")
+	assert.Equal(t, []schema.ChatMessage{
+		schema.AIChatMessage{Text: "foo"},
+		schema.HumanChatMessage{Text: "bar"},
+	}, h.Messages())
+
+	h = NewChatMessageHistory(
+		WithPreviousMessages([]schema.ChatMessage{
+			schema.AIChatMessage{Text: "foo"},
+			schema.SystemChatMessage{Text: "bar"},
+		}),
+	)
+	h.AddUserMessage("zoo")
+	assert.Equal(t, []schema.ChatMessage{
+		schema.AIChatMessage{Text: "foo"},
+		schema.SystemChatMessage{Text: "bar"},
+		schema.HumanChatMessage{Text: "zoo"},
+	}, h.Messages())
+}