From c9480116aab43405c7d3a7c2d4bf3441131b31d6 Mon Sep 17 00:00:00 2001 From: zivkovicn Date: Mon, 24 Jul 2023 18:49:38 +0200 Subject: [PATCH] feature-adding-max-token-size-memory | logic update --- memory/token_buffer.go | 13 ++++--------- outputparser/simple.go | 15 +++++++++++---- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/memory/token_buffer.go b/memory/token_buffer.go index fcd9b1b2d..b150b02c9 100644 --- a/memory/token_buffer.go +++ b/memory/token_buffer.go @@ -68,15 +68,10 @@ func (tb *TokenBuffer) Clear() error { } func (tb *TokenBuffer) getNumTokensFromMessages() (int, error) { - sum := 0 - for _, message := range tb.ChatHistory.Messages() { - bufferString, err := schema.GetBufferString([]schema.ChatMessage{message}, tb.Buffer.HumanPrefix, tb.Buffer.AIPrefix) - if err != nil { - return 0, err - } - - sum += tb.LLM.GetNumTokens(bufferString) + bufferString, err := schema.GetBufferString(tb.ChatHistory.Messages(), tb.Buffer.HumanPrefix, tb.Buffer.AIPrefix) + if err != nil { + return 0, err } - return sum, nil + return tb.LLM.GetNumTokens(bufferString), nil } diff --git a/outputparser/simple.go b/outputparser/simple.go index 248cfa1b5..412980d7f 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" }