feature-adding-max-token-size-memory | logic update

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
 }

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