From 735bab5e32a49a9f9f5f12142208f41be74dcd6d Mon Sep 17 00:00:00 2001
From: mudler <mudler@localai.io>
Date: Wed, 18 Dec 2024 12:41:18 +0100
Subject: [PATCH] Fixups to long memory, and avoid re-initing RAGDB recursively

Signed-off-by: mudler <mudler@localai.io>
---
 agent/agent.go        | 24 +++++++++++++-----------
 example/webui/main.go |  2 +-
 example/webui/rag.go  |  4 ++--
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 5057d9f..40d285b 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -282,8 +282,6 @@ func (a *Agent) consumeJob(job *Job, role string) {
 	// We are self evaluating if we consume the job as a system role
 	selfEvaluation := role == SystemRole
 
-	memory := a.options.enableKB && a.options.ragdb != nil
-
 	a.Lock()
 	// Set the action context
 	ctx, cancel := context.WithCancel(context.Background())
@@ -337,18 +335,27 @@ func (a *Agent) consumeJob(job *Job, role string) {
 		}
 	}
 
+	if job.Text != "" {
+		a.currentConversation = append(a.currentConversation, openai.ChatCompletionMessage{
+			Role:    role,
+			Content: job.Text,
+		})
+	}
+
 	// TODO: move to a promptblock?
 	// RAG
-	if memory {
+	if a.options.enableLongTermMemory && len(a.currentConversation) > 0 {
 		// Walk conversation from bottom to top, and find the first message of the user
 		// to use it as a query to the KB
 		var userMessage string
-		for i := len(a.currentConversation) - 1; i >= 0; i-- {
+		for i := len(a.currentConversation); i == 0; i-- {
+			xlog.Info("[Long term memory] Conversation", "role", a.currentConversation[i].Role, "Content", a.currentConversation[i].Content)
 			if a.currentConversation[i].Role == "user" {
 				userMessage = a.currentConversation[i].Content
 				break
 			}
 		}
+		xlog.Info("[Long term memory] User message", "agent", a.Character.Name, "message", userMessage)
 
 		if userMessage != "" {
 			results, err := a.options.ragdb.Search(userMessage, a.options.kbResults)
@@ -380,13 +387,8 @@ func (a *Agent) consumeJob(job *Job, role string) {
 					}}, a.currentConversation...)
 			}
 		}
-	}
-
-	if job.Text != "" {
-		a.currentConversation = append(a.currentConversation, openai.ChatCompletionMessage{
-			Role:    role,
-			Content: job.Text,
-		})
+	} else {
+		xlog.Info("[Long term memory] No conversation available", "agent", a.Character.Name)
 	}
 
 	var pickTemplate string
diff --git a/example/webui/main.go b/example/webui/main.go
index 642bd14..6b21d1a 100644
--- a/example/webui/main.go
+++ b/example/webui/main.go
@@ -82,7 +82,7 @@ func main() {
 
 	if len(db.Database) > 0 && kbdisableIndexing != "true" {
 		xlog.Info("Loading knowledgebase from disk, to skip run with KBDISABLEINDEX=true")
-		if err := db.SaveAllToStore(); err != nil {
+		if err := db.PopulateRAGDB(); err != nil {
 			xlog.Info("Error storing in the KB", err)
 		}
 	}
diff --git a/example/webui/rag.go b/example/webui/rag.go
index a8c27b0..30bbeae 100644
--- a/example/webui/rag.go
+++ b/example/webui/rag.go
@@ -63,13 +63,13 @@ func NewInMemoryDB(knowledgebase string, store RAGDB) (*InMemoryDatabase, error)
 	}, nil
 }
 
-func (db *InMemoryDatabase) SaveAllToStore() error {
+func (db *InMemoryDatabase) PopulateRAGDB() error {
 	for _, d := range db.Database {
 		if d == "" {
 			// skip empty chunks
 			continue
 		}
-		err := db.Store(d)
+		err := db.RAGDB.Store(d)
 		if err != nil {
 			return fmt.Errorf("error storing in the KB: %w", err)
 		}