From 351aa80b7433d6654ab05f4eb8216707d5f060c9 Mon Sep 17 00:00:00 2001
From: mudler <mudler@localai.io>
Date: Tue, 17 Dec 2024 22:24:38 +0100
Subject: [PATCH] Use InMemoryDB in place of a RAGDB to store things as files

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

diff --git a/example/webui/main.go b/example/webui/main.go
index 6e73b85..642bd14 100644
--- a/example/webui/main.go
+++ b/example/webui/main.go
@@ -55,34 +55,34 @@ func main() {
 	stateDir := cwd + "/pool"
 	os.MkdirAll(stateDir, 0755)
 
-	var dbStore RAGDB
+	var ragDB RAGDB
 	lai := llm.NewClient(apiKey, apiURL+"/v1", timeout)
 
 	switch vectorStore {
 	case "localai":
 		laiStore := rag.NewStoreClient(apiURL, apiKey)
-		dbStore = rag.NewLocalAIRAGDB(laiStore, lai)
+		ragDB = rag.NewLocalAIRAGDB(laiStore, lai)
 	default:
 		var err error
-		dbStore, err = rag.NewChromemDB("local-agent-framework", stateDir, lai, embeddingModel)
+		ragDB, err = rag.NewChromemDB("local-agent-framework", stateDir, lai, embeddingModel)
 		if err != nil {
 			panic(err)
 		}
 	}
 
-	pool, err := NewAgentPool(testModel, apiURL, stateDir, dbStore)
+	db, err := NewInMemoryDB(stateDir, ragDB)
 	if err != nil {
 		panic(err)
 	}
 
-	db, err := NewInMemoryDB(stateDir, dbStore)
+	pool, err := NewAgentPool(testModel, apiURL, stateDir, db)
 	if err != nil {
 		panic(err)
 	}
 
 	if len(db.Database) > 0 && kbdisableIndexing != "true" {
 		xlog.Info("Loading knowledgebase from disk, to skip run with KBDISABLEINDEX=true")
-		if err := db.SaveToStore(); err != nil {
+		if err := db.SaveAllToStore(); err != nil {
 			xlog.Info("Error storing in the KB", err)
 		}
 	}
diff --git a/example/webui/rag.go b/example/webui/rag.go
index 54c3283..b71453b 100644
--- a/example/webui/rag.go
+++ b/example/webui/rag.go
@@ -20,10 +20,10 @@ import (
 )
 
 type InMemoryDatabase struct {
+	RAGDB
 	sync.Mutex
 	Database []string
 	path     string
-	rag      RAGDB
 }
 
 func loadDB(path string) ([]string, error) {
@@ -48,7 +48,7 @@ func NewInMemoryDB(knowledgebase string, store RAGDB) (*InMemoryDatabase, error)
 		return &InMemoryDatabase{
 			Database: []string{},
 			path:     poolfile,
-			rag:      store,
+			RAGDB:    store,
 		}, nil
 	}
 
@@ -57,21 +57,21 @@ func NewInMemoryDB(knowledgebase string, store RAGDB) (*InMemoryDatabase, error)
 		return nil, err
 	}
 	return &InMemoryDatabase{
+		RAGDB:    store,
 		Database: poolData,
 		path:     poolfile,
-		rag:      store,
 	}, nil
 }
 
-func (db *InMemoryDatabase) SaveToStore() error {
+func (db *InMemoryDatabase) SaveAllToStore() error {
 	for _, d := range db.Database {
 		if d == "" {
 			// skip empty chunks
 			continue
 		}
-		err := db.rag.Store(d)
+		err := db.Store(d)
 		if err != nil {
-			return fmt.Errorf("Error storing in the KB: %w", err)
+			return fmt.Errorf("error storing in the KB: %w", err)
 		}
 	}
 
@@ -82,17 +82,17 @@ func (db *InMemoryDatabase) Reset() error {
 	db.Lock()
 	db.Database = []string{}
 	db.Unlock()
-	if err := db.rag.Reset(); err != nil {
+	if err := db.RAGDB.Reset(); err != nil {
 		return err
 	}
 	return db.SaveDB()
 }
 
-func (db *InMemoryDatabase) AddEntry(entry string) error {
+func (db *InMemoryDatabase) Store(entry string) error {
 	db.Lock()
 	defer db.Unlock()
 	db.Database = append(db.Database, entry)
-	return nil
+	return db.RAGDB.Store(entry)
 }
 
 func (db *InMemoryDatabase) SaveDB() error {
@@ -150,15 +150,11 @@ func StringsToKB(db *InMemoryDatabase, chunkSize int, content ...string) {
 		xlog.Info("chunks: ", len(chunks))
 		for _, chunk := range chunks {
 			xlog.Info("Chunk size: ", len(chunk))
-			db.AddEntry(chunk)
+			db.Store(chunk)
 		}
 
 		db.SaveDB()
 	}
-
-	if err := db.SaveToStore(); err != nil {
-		xlog.Info("Error storing in the KB", err)
-	}
 }
 
 // splitParagraphIntoChunks takes a paragraph and a maxChunkSize as input,