lock during store and replace

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Author: mudler

rag/collection.go

@@ -52,7 +52,7 @@ func NewPersistentLocalAICollection(llmClient *openai.Client, apiURL, apiKey, co
 
 	// TODO: This does not work as we do not have .Reset().
 	// The problem is that LocalAI stores are not persistent either and do not allow upserts.
-	persistentKB.repopulate()
+	persistentKB.Repopulate()
 
 	return persistentKB
 }

rag/persistency.go

@@ -120,11 +120,16 @@ func (db *PersistentKB) save() error {
 	return os.WriteFile(db.path, data, 0644)
 }
 
-// repopulate reinitializes the persistent knowledge base with the files that were added to it.
-func (db *PersistentKB) repopulate() error {
+func (db *PersistentKB) Count() int {
 	db.Lock()
 	defer db.Unlock()
 
+	return db.Engine.Count()
+}
+
+// repopulate reinitializes the persistent knowledge base with the files that were added to it.
+func (db *PersistentKB) repopulate() error {
+
 	if err := db.Engine.Reset(); err != nil {
 		return fmt.Errorf("failed to reset engine: %w", err)
 	}
@@ -141,6 +146,13 @@ func (db *PersistentKB) repopulate() error {
 	return nil
 }
 
+func (db *PersistentKB) Repopulate() error {
+	db.Lock()
+	defer db.Unlock()
+
+	return db.repopulate()
+}
+
 // Store stores an entry in the persistent knowledge base.
 func (db *PersistentKB) ListDocuments() []string {
 	db.Lock()
@@ -173,6 +185,10 @@ func (db *PersistentKB) Store(entry string, metadata map[string]string) error {
 	db.Lock()
 	defer db.Unlock()
 
+	return db.storeFile(entry, metadata)
+}
+
+func (db *PersistentKB) storeFile(entry string, metadata map[string]string) error {
 	fileName := filepath.Base(entry)
 
 	// copy file to assetDir (if it's a file)
@@ -193,19 +209,17 @@ func (db *PersistentKB) Store(entry string, metadata map[string]string) error {
 }
 
 func (db *PersistentKB) StoreOrReplace(entry string, metadata map[string]string) error {
-	db.Lock()
 	fileName := filepath.Base(entry)
 	_, ok := db.index[fileName]
-	db.Unlock()
 	// Check if we have it already in the index
 	if ok {
 		xlog.Info("Data already exists for entry", "entry", entry, "index", db.index)
-		if err := db.RemoveEntry(fileName); err != nil {
+		if err := db.removeFileEntry(fileName); err != nil {
 			return fmt.Errorf("failed to remove entry: %w", err)
 		}
 	}
 
-	return db.Store(entry, metadata)
+	return db.storeFile(entry, metadata)
 }
 
 func (db *PersistentKB) store(metadata map[string]string, files ...string) ([]engine.Result, error) {
@@ -231,8 +245,14 @@ func (db *PersistentKB) store(metadata map[string]string, files ...string) ([]en
 	return results, nil
 }
 
-// RemoveEntry removes an entry from the persistent knowledge base.
 func (db *PersistentKB) RemoveEntry(entry string) error {
+	db.Lock()
+	defer db.Unlock()
+
+	return db.removeFileEntry(entry)
+}
+
+func (db *PersistentKB) removeFileEntry(entry string) error {
 
 	xlog.Info("Removing entry", "entry", entry)
 	if os.Getenv("LOCALRECALL_REPOPULATE_DELETE") != "true" {
@@ -261,25 +281,20 @@ func (db *PersistentKB) RemoveEntry(entry string) error {
 			}
 		}
 
-		db.Lock()
-
 		xlog.Info("Deleting entry from index", "entry", entry)
 		delete(db.index, entry)
 
 		xlog.Info("Removing entry from disk", "file", e)
 		os.Remove(e)
-		db.Unlock()
 		return db.save()
 	}
 
-	db.Lock()
 	for e := range db.index {
 		if e == entry {
 			os.Remove(filepath.Join(db.assetDir, e))
 			break
 		}
 	}
-	db.Unlock()
 
 	// TODO: this is suboptimal, but currently chromem does not support deleting single entities
 	return db.repopulate()

test/e2e/common_test.go

@@ -20,7 +20,7 @@ const (
 	DefaultUpdateInterval = time.Hour
 
 	// TestTimeout is the default timeout for Eventually blocks
-	TestTimeout = 5 * time.Second
+	TestTimeout = 1 * time.Minute
 
 	// TestPollingInterval is the default polling interval for Eventually blocks
 	TestPollingInterval = 500 * time.Millisecond

test/e2e/source_manager_test.go

@@ -247,16 +247,19 @@ var _ = Describe("SourceManager", func() {
 				return kb.ListDocuments()
 			}, 2*time.Minute, 5*time.Second).Should(HaveLen(1))
 
+			// Wait for initial content to be fetched
+			Eventually(func() int {
+				return kb.Count()
+			}, 2*time.Minute, 5*time.Second).Should(Equal(25))
+
 			// Let it run for 2 minutes and check for duplicates
 			Consistently(func() int {
-				e, ok := kb.Engine.(*engine.ChromemDB)
-				Expect(ok).To(BeTrue())
-				return e.Count()
+				return kb.Count()
 			}, 3*time.Minute, 5*time.Second).Should(Equal(25))
 
 			// Verify that search results don't contain duplicates
 			Consistently(func() bool {
-				results, err := kb.Engine.Search("What is the Black-crowned barwing?", 3)
+				results, err := kb.Engine.Search("What is the Black-crowned barwing?", 5)
 				if err != nil {
 					fmt.Println("Error searching for content", err)
 					return false