Always store in KB when adding new content

Signed-off-by: mudler <mudler@localai.io>
This commit is contained in:
mudler
2024-12-18 11:29:08 +01:00
parent 351aa80b74
commit c2035c95c5

View File

@@ -88,23 +88,29 @@ func (db *InMemoryDatabase) Reset() error {
return db.SaveDB() return db.SaveDB()
} }
func (db *InMemoryDatabase) Store(entry string) error { func (db *InMemoryDatabase) save() error {
db.Lock()
defer db.Unlock()
db.Database = append(db.Database, entry)
return db.RAGDB.Store(entry)
}
func (db *InMemoryDatabase) SaveDB() error {
db.Lock()
defer db.Unlock()
data, err := json.Marshal(db.Database) data, err := json.Marshal(db.Database)
if err != nil { if err != nil {
return err return err
} }
err = os.WriteFile(db.path, data, 0644) return os.WriteFile(db.path, data, 0644)
return err }
func (db *InMemoryDatabase) Store(entry string) error {
db.Lock()
defer db.Unlock()
db.Database = append(db.Database, entry)
if err := db.RAGDB.Store(entry); err != nil {
return err
}
return db.save()
}
func (db *InMemoryDatabase) SaveDB() error {
db.Lock()
defer db.Unlock()
return db.save()
} }
func getWebPage(url string) (string, error) { func getWebPage(url string) (string, error) {
@@ -121,7 +127,7 @@ func getWebPage(url string) (string, error) {
return html2text.FromString(string(body), html2text.Options{PrettyTables: true}) return html2text.FromString(string(body), html2text.Options{PrettyTables: true})
} }
func Sitemap(url string) (res []string, err error) { func getWebSitemap(url string) (res []string, err error) {
err = sitemap.ParseFromSite(url, func(e sitemap.Entry) error { err = sitemap.ParseFromSite(url, func(e sitemap.Entry) error {
xlog.Info("Sitemap page: " + e.GetLocation()) xlog.Info("Sitemap page: " + e.GetLocation())
content, err := getWebPage(e.GetLocation()) content, err := getWebPage(e.GetLocation())
@@ -134,7 +140,7 @@ func Sitemap(url string) (res []string, err error) {
} }
func WebsiteToKB(website string, chunkSize int, db *InMemoryDatabase) { func WebsiteToKB(website string, chunkSize int, db *InMemoryDatabase) {
content, err := Sitemap(website) content, err := getWebSitemap(website)
if err != nil { if err != nil {
xlog.Info("Error walking sitemap for website", err) xlog.Info("Error walking sitemap for website", err)
} }
@@ -152,8 +158,6 @@ func StringsToKB(db *InMemoryDatabase, chunkSize int, content ...string) {
xlog.Info("Chunk size: ", len(chunk)) xlog.Info("Chunk size: ", len(chunk))
db.Store(chunk) db.Store(chunk)
} }
db.SaveDB()
} }
} }