Repository: tmc/langchaingo
Pull request #217 (2023-07-24 16:52:45)

memory/tokenbuffer: adding max size token and prune of the messages
### 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 a new buffer `TokenBuffer` which can take the max number of tokens, and use the token count for the given context to prune the old history of the messages in favor of new ones and to not hit the limit.
- fixing issue with simple parser, to trim the left whitespace, when we get response


Patch:
diff --git a/chains/chains.go b/chains/chains.go
index 56c14f16..15649840 100644
--- a/chains/chains.go
+++ b/chains/chains.go
@@ -19,9 +19,9 @@ type Chain interface {
 	Call(ctx context.Context, inputs map[string]any, options ...ChainCallOption) (map[string]any, error)
 	// GetMemory gets the memory of the chain.
 	GetMemory() schema.Memory
-	// InputKeys returns the input keys the chain expects.
+	// GetInputKeys returns the input keys the chain expects.
 	GetInputKeys() []string
-	// OutputKeys returns the output keys the chain returns.
+	// GetOutputKeys returns the output keys the chain returns.
 	GetOutputKeys() []string
 }
 
diff --git a/memory/chat.go b/memory/chat.go
index dd2165d5..1f5798f4 100644
--- a/memory/chat.go
+++ b/memory/chat.go
@@ -37,3 +37,7 @@ func (h *ChatMessageHistory) Clear() {
 func (h *ChatMessageHistory) AddMessage(message schema.ChatMessage) {
 	h.messages = append(h.messages, message)
 }
+
+func (h *ChatMessageHistory) SetMessages(messages []schema.ChatMessage) {
+	h.messages = messages
+}
diff --git a/memory/token_buffer.go b/memory/token_buffer.go
new file mode 100644
index 00000000..6244a7c5
--- /dev/null
+++ b/memory/token_buffer.go
@@ -0,0 +1,81 @@
+package memory
+
+import (
+	"github.com/tmc/langchaingo/llms"
+	"github.com/tmc/langchaingo/schema"
+)
+
+// TokenBuffer for storing conversation memory.
+type TokenBuffer struct {
+	Buffer
+	LLM           llms.LanguageModel
+	MaxTokenLimit int
+}
+
+// Statically assert that TokenBuffer implement the memory interface.
+var _ schema.Memory = &TokenBuffer{}
+
+// 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...),
+	}
+
+	return tb
+}
+
+// MemoryVariables uses Buffer method for memory variables.
+func (tb *TokenBuffer) MemoryVariables() []string {
+	return tb.Buffer.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)
+}
+
+// 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)
+	if err != nil {
+		return err
+	}
+	currBufferLength, err := tb.getNumTokensFromMessages()
+	if err != nil {
+		return err
+	}
+
+	if currBufferLength > tb.MaxTokenLimit {
+		// 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 {
+				break
+			}
+
+			tb.ChatHistory.SetMessages(append(tb.ChatHistory.Messages()[:0], tb.ChatHistory.Messages()[1:]...))
+			currBufferLength, err = tb.getNumTokensFromMessages()
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+// Clear uses Buffer method for clearing buffer memory.
+func (tb *TokenBuffer) Clear() error {
+	return tb.Buffer.Clear()
+}
+
+func (tb *TokenBuffer) getNumTokensFromMessages() (int, error) {
+	bufferString, err := schema.GetBufferString(tb.ChatHistory.Messages(), tb.Buffer.HumanPrefix, tb.Buffer.AIPrefix)
+	if err != nil {
+		return 0, err
+	}
+
+	return tb.LLM.GetNumTokens(bufferString), nil
+}
diff --git a/outputparser/simple.go b/outputparser/simple.go
index 248cfa1b..412980d7 100644
--- a/outputparser/simple.go
+++ b/outputparser/simple.go
@@ -1,6 +1,10 @@
 package outputparser
 
-import "github.com/tmc/langchaingo/schema"
+import (
+	"strings"
+
+	"github.com/tmc/langchaingo/schema"
+)
 
 // Simple is an output parser that does nothing.
 type Simple struct{}
@@ -9,9 +13,12 @@ func NewSimple() Simple { return Simple{} }
 
 var _ schema.OutputParser[any] = Simple{}
 
-func (p Simple) GetFormatInstructions() string  { return "" }
-func (p Simple) Parse(text string) (any, error) { return text, nil }
+func (p Simple) GetFormatInstructions() string { return "" }
+func (p Simple) Parse(text string) (any, error) {
+	return strings.TrimSpace(text), nil
+}
+
 func (p Simple) ParseWithPrompt(text string, _ schema.PromptValue) (any, error) {
-	return text, nil
+	return strings.TrimSpace(text), nil
 }
 func (p Simple) Type() string { return "simple_parser" }
diff --git a/schema/chat_message_history.go b/schema/chat_message_history.go
index 1fdde7a0..2507ee42 100644
--- a/schema/chat_message_history.go
+++ b/schema/chat_message_history.go
@@ -16,4 +16,7 @@ type ChatMessageHistory interface {
 
 	// Messages get all messages from the store
 	Messages() []ChatMessage
+
+	// SetMessages replaces existing messages in the store
+	SetMessages(messages []ChatMessage)
 }
diff --git a/schema/memory.go b/schema/memory.go
index ebc00983..a6c75e2b 100644
--- a/schema/memory.go
+++ b/schema/memory.go
@@ -2,12 +2,12 @@ package schema
 
 // Memory is the interface for memory in chains.
 type Memory interface {
-	// Input keys this memory class will load dynamically.
+	// MemoryVariables Input keys this memory class will load dynamically.
 	MemoryVariables() []string
-	// Return key-value pairs given the text input to the chain.
+	// LoadMemoryVariables Return key-value pairs given the text input to the chain.
 	// If None, return all memories
 	LoadMemoryVariables(inputs map[string]any) (map[string]any, error)
-	// Save the context of this model run to memory.
+	// SaveContext Save the context of this model run to memory.
 	SaveContext(inputs map[string]any, outputs map[string]any) error
 	// Clear memory contents.
 	Clear() error


Test patch:
diff --git a/chains/conversation_test.go b/chains/conversation_test.go
index 87409fbc..cd36960e 100644
--- a/chains/conversation_test.go
+++ b/chains/conversation_test.go
@@ -13,13 +13,14 @@ import (
 
 func TestConversation(t *testing.T) {
 	t.Parallel()
+
 	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
 		t.Skip("OPENAI_API_KEY not set")
 	}
-	model, err := openai.New()
+	llm, err := openai.New()
 	require.NoError(t, err)
 
-	c := NewConversation(model, memory.NewBuffer())
+	c := NewConversation(llm, memory.NewBuffer())
 	_, err = Run(context.Background(), c, "Hi! I'm Jim")
 	require.NoError(t, err)
 
@@ -27,3 +28,27 @@ func TestConversation(t *testing.T) {
 	require.NoError(t, err)
 	require.True(t, strings.Contains(res, "Jim"), `result does not contain the keyword 'Jim'`)
 }
+
+func TestConversationMemoryPrune(t *testing.T) {
+	t.Parallel()
+
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+
+	llm, err := openai.New()
+	require.NoError(t, err)
+
+	c := NewConversation(llm, memory.NewTokenBuffer(llm, 50))
+	_, err = Run(context.Background(), c, "Hi! I'm Jim")
+	require.NoError(t, err)
+
+	res, err := Run(context.Background(), c, "What is my name?")
+	require.NoError(t, err)
+	require.True(t, strings.Contains(res, "Jim"), `result does contain the keyword 'Jim'`)
+
+	// this message will hit the maxTokenLimit and will initiate the prune of the messages to fit the context
+	res, err = Run(context.Background(), c, "Are you sure that my name is Jim?")
+	require.NoError(t, err)
+	require.True(t, strings.Contains(res, "Jim"), `result does contain the keyword 'Jim'`)
+}
diff --git a/chains/map_reduce_test.go b/chains/map_reduce_test.go
index 7bec3c48..8fba5d27 100644
--- a/chains/map_reduce_test.go
+++ b/chains/map_reduce_test.go
@@ -54,5 +54,5 @@ func TestMapReduce(t *testing.T) {
 		{PageContent: "doo"},
 	})
 	require.NoError(t, err)
-	require.Equal(t, "foo\n\nboo\n\nzoo\n\ndoo\n\n", result)
+	require.Equal(t, "foo\n\nboo\n\nzoo\n\ndoo", result)
 }
diff --git a/memory/buffer_test.go b/memory/buffer_test.go
index a9a81931..3a6c2b41 100644
--- a/memory/buffer_test.go
+++ b/memory/buffer_test.go
@@ -86,6 +86,9 @@ func (t testChatMessageHistory) AddMessage(_ schema.ChatMessage) {
 func (t testChatMessageHistory) Clear() {
 }
 
+func (t testChatMessageHistory) SetMessages(_ []schema.ChatMessage) {
+}
+
 func (t testChatMessageHistory) Messages() []schema.ChatMessage {
 	return []schema.ChatMessage{
 		schema.HumanChatMessage{Content: "user message test"},
diff --git a/memory/token_buffer_test.go b/memory/token_buffer_test.go
new file mode 100644
index 00000000..4ea4b2f2
--- /dev/null
+++ b/memory/token_buffer_test.go
@@ -0,0 +1,93 @@
+package memory
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/tmc/langchaingo/llms/openai"
+	"github.com/tmc/langchaingo/schema"
+)
+
+func TestTokenBufferMemory(t *testing.T) {
+	t.Parallel()
+
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+
+	llm, err := openai.New()
+	require.NoError(t, err)
+	m := NewTokenBuffer(llm, 2000)
+
+	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 TestTokenBufferMemoryReturnMessage(t *testing.T) {
+	t.Parallel()
+
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+
+	llm, err := openai.New()
+	require.NoError(t, err)
+	m := NewTokenBuffer(llm, 2000, WithReturnMessages(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{Content: "bar"},
+			schema.AIChatMessage{Content: "foo"},
+		}),
+	)
+
+	expected2 := map[string]any{"history": expectedChatHistory.Messages()}
+	assert.Equal(t, expected2, result2)
+}
+
+func TestTokenBufferMemoryWithPreLoadedHistory(t *testing.T) {
+	t.Parallel()
+
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+
+	llm, err := openai.New()
+	require.NoError(t, err)
+
+	m := NewTokenBuffer(llm, 2000, WithChatHistory(NewChatMessageHistory(
+		WithPreviousMessages([]schema.ChatMessage{
+			schema.HumanChatMessage{Content: "bar"},
+			schema.AIChatMessage{Content: "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)
+}