feat: support separate knowledge bases for each agent

Also allow to export/import KB

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

agent/agent.go

@@ -45,6 +45,7 @@ type RAGDB interface {
 	Store(s string) error
 	Reset() error
 	Search(s string, similarEntries int) ([]string, error)
+	Count() int
 }
 
 func New(opts ...Option) (*Agent, error) {
@@ -235,6 +236,10 @@ func (a *Agent) Paused() bool {
 	return a.pause
 }
 
+func (a *Agent) Memory() RAGDB {
+	return a.options.ragdb
+}
+
 func (a *Agent) runAction(chosenAction Action, params action.ActionParams) (result string, err error) {
 	for _, action := range a.systemInternalActions() {
 		if action.Definition().Name == chosenAction.Definition().Name {

example/webui/agentpool.go

@@ -24,6 +24,7 @@ type AgentPool struct {
 	agents        map[string]*Agent
 	managers      map[string]Manager
 	agentStatus   map[string]*Status
+	agentMemory   map[string]*InMemoryDatabase
 	apiURL, model string
 	ragDB         RAGDB
 }
@@ -76,6 +77,7 @@ func NewAgentPool(model, apiURL, directory string, RagDB RAGDB) (*AgentPool, err
 			pool:        make(map[string]AgentConfig),
 			agentStatus: make(map[string]*Status),
 			managers:    make(map[string]Manager),
+			agentMemory: make(map[string]*InMemoryDatabase),
 		}, nil
 	}
 
@@ -92,6 +94,7 @@ func NewAgentPool(model, apiURL, directory string, RagDB RAGDB) (*AgentPool, err
 		agents:      make(map[string]*Agent),
 		managers:    make(map[string]Manager),
 		agentStatus: map[string]*Status{},
+		agentMemory: map[string]*InMemoryDatabase{},
 		pool:        *poolData,
 	}, nil
 }
@@ -144,7 +147,14 @@ func (a *AgentPool) startAgentWithConfig(name string, config *AgentConfig) error
 
 	actions := config.availableActions(ctx)
 
-	stateFile, characterFile := a.stateFiles(name)
+	stateFile, characterFile, knowledgeBase := a.stateFiles(name)
+
+	agentDB, err := NewInMemoryDB(knowledgeBase, a.ragDB)
+	if err != nil {
+		return err
+	}
+
+	a.agentMemory[name] = agentDB
 
 	actionsLog := []string{}
 	for _, action := range actions {
@@ -179,7 +189,7 @@ func (a *AgentPool) startAgentWithConfig(name string, config *AgentConfig) error
 		WithStateFile(stateFile),
 		WithCharacterFile(characterFile),
 		WithTimeout(timeout),
-		WithRAGDB(a.ragDB),
+		WithRAGDB(agentDB),
 		WithAgentReasoningCallback(func(state ActionCurrentState) bool {
 			xlog.Info(
 				"Agent is thinking",
@@ -355,20 +365,22 @@ func (a *AgentPool) Start(name string) error {
 	return fmt.Errorf("agent %s not found", name)
 }
 
-func (a *AgentPool) stateFiles(name string) (string, string) {
+func (a *AgentPool) stateFiles(name string) (string, string, string) {
 	stateFile := filepath.Join(a.pooldir, fmt.Sprintf("%s.state.json", name))
 	characterFile := filepath.Join(a.pooldir, fmt.Sprintf("%s.character.json", name))
+	knowledgeBaseFile := filepath.Join(a.pooldir, fmt.Sprintf("%s.knowledgebase.json", name))
 
-	return stateFile, characterFile
+	return stateFile, characterFile, knowledgeBaseFile
 }
 
 func (a *AgentPool) Remove(name string) error {
 
 	// Cleanup character and state
-	stateFile, characterFile := a.stateFiles(name)
+	stateFile, characterFile, knowledgeBaseFile := a.stateFiles(name)
 
 	os.Remove(stateFile)
 	os.Remove(characterFile)
+	os.Remove(knowledgeBaseFile)
 
 	a.Stop(name)
 	delete(a.agents, name)
@@ -391,6 +403,10 @@ func (a *AgentPool) GetAgent(name string) *Agent {
 	return a.agents[name]
 }
 
+func (a *AgentPool) GetAgentMemory(name string) *InMemoryDatabase {
+	return a.agentMemory[name]
+}
+
 func (a *AgentPool) GetConfig(name string) *AgentConfig {
 	agent, exists := a.pool[name]
 	if !exists {

example/webui/app.go

@@ -24,15 +24,72 @@ type (
 	}
 )
 
-func (a *App) KnowledgeBaseReset(db *InMemoryDatabase) func(c *fiber.Ctx) error {
+func (a *App) KnowledgeBaseReset(pool *AgentPool) func(c *fiber.Ctx) error {
 	return func(c *fiber.Ctx) error {
+		db := pool.GetAgentMemory(c.Params("name"))
 		db.Reset()
-		return c.Redirect("/knowledgebase")
+		return c.Redirect("/knowledgebase/" + c.Params("name"))
 	}
 }
 
-func (a *App) KnowledgeBaseFile(db *InMemoryDatabase) func(c *fiber.Ctx) error {
+func (a *App) KnowledgeBaseExport(pool *AgentPool) func(c *fiber.Ctx) error {
 	return func(c *fiber.Ctx) error {
+		db := pool.GetAgentMemory(c.Params("name"))
+		knowledgeBase := db.Data()
+
+		c.Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s.knowledgebase.json", c.Params("name")))
+		return c.JSON(knowledgeBase)
+	}
+}
+
+func (a *App) KnowledgeBaseImport(pool *AgentPool) func(c *fiber.Ctx) error {
+	return func(c *fiber.Ctx) error {
+		file, err := c.FormFile("file")
+		if err != nil {
+			// Handle error
+			return err
+		}
+
+		os.MkdirAll("./uploads", os.ModePerm)
+
+		destination := fmt.Sprintf("./uploads/%s", file.Filename)
+		if err := c.SaveFile(file, destination); err != nil {
+			// Handle error
+			return err
+		}
+
+		data, err := os.ReadFile(destination)
+		if err != nil {
+			return err
+		}
+
+		knowledge := []string{}
+		if err := json.Unmarshal(data, &knowledge); err != nil {
+			return err
+		}
+
+		if len(knowledge) > 0 {
+			xlog.Info("Importing agent KB")
+			db := pool.GetAgentMemory(c.Params("name"))
+			db.Reset()
+
+			for _, k := range knowledge {
+				db.Store(k)
+			}
+
+		} else {
+			return fmt.Errorf("Empty knowledge base")
+		}
+
+		return c.Redirect("/agents")
+	}
+}
+
+func (a *App) KnowledgeBaseFile(pool *AgentPool) func(c *fiber.Ctx) error {
+	return func(c *fiber.Ctx) error {
+		agent := pool.GetAgent(c.Params("name"))
+		db := agent.Memory()
+
 		// https://golang.withcodeexample.com/blog/file-upload-handling-golang-fiber-guide/
 		file, err := c.FormFile("file")
 		if err != nil {
@@ -78,8 +135,11 @@ func (a *App) KnowledgeBaseFile(db *InMemoryDatabase) func(c *fiber.Ctx) error {
 	}
 }
 
-func (a *App) KnowledgeBase(db *InMemoryDatabase) func(c *fiber.Ctx) error {
+func (a *App) KnowledgeBase(pool *AgentPool) func(c *fiber.Ctx) error {
 	return func(c *fiber.Ctx) error {
+		agent := pool.GetAgent(c.Params("name"))
+		db := agent.Memory()
+
 		payload := struct {
 			URL       string `form:"url"`
 			ChunkSize int    `form:"chunk_size"`
@@ -100,7 +160,7 @@ func (a *App) KnowledgeBase(db *InMemoryDatabase) func(c *fiber.Ctx) error {
 
 		go WebsiteToKB(website, chunkSize, db)
 
-		return c.Redirect("/knowledgebase")
+		return c.Redirect("/knowledgebase/" + c.Params("name"))
 	}
 }
 

example/webui/main.go

@@ -6,8 +6,6 @@ import (
 	"net/http"
 	"os"
 
-	"github.com/mudler/local-agent-framework/xlog"
-
 	"github.com/donseba/go-htmx"
 	fiber "github.com/gofiber/fiber/v2"
 	"github.com/gofiber/template/html/v2"
@@ -21,7 +19,6 @@ var testModel = os.Getenv("TEST_MODEL")
 var apiURL = os.Getenv("API_URL")
 var apiKey = os.Getenv("API_KEY")
 var vectorStore = os.Getenv("VECTOR_STORE")
-var kbdisableIndexing = os.Getenv("KBDISABLEINDEX")
 var timeout = os.Getenv("TIMEOUT")
 var embeddingModel = os.Getenv("EMBEDDING_MODEL")
 
@@ -70,23 +67,11 @@ func main() {
 		}
 	}
 
-	db, err := NewInMemoryDB(stateDir, ragDB)
+	pool, err := NewAgentPool(testModel, apiURL, stateDir, ragDB)
 	if err != nil {
 		panic(err)
 	}
 
-	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.PopulateRAGDB(); err != nil {
-			xlog.Info("Error storing in the KB", err)
-		}
-	}
-
 	app := &App{
 		htmx: htmx.New(),
 		pool: pool,
@@ -102,7 +87,7 @@ func main() {
 		Views: engine,
 	})
 
-	RegisterRoutes(webapp, pool, db, app)
+	RegisterRoutes(webapp, pool, app)
 
 	log.Fatal(webapp.Listen(":3000"))
 }

example/webui/rag.go

@@ -9,7 +9,6 @@ import (
 
 	"net/http"
 	"os"
-	"path/filepath"
 	"strings"
 	"sync"
 
@@ -37,12 +36,10 @@ func loadDB(path string) ([]string, error) {
 	return poolData, err
 }
 
-func NewInMemoryDB(knowledgebase string, store RAGDB) (*InMemoryDatabase, error) {
+func NewInMemoryDB(poolfile string, store RAGDB) (*InMemoryDatabase, error) {
 	// if file exists, try to load an existing pool.
 	// if file does not exist, create a new pool.
 
-	poolfile := filepath.Join(knowledgebase, "knowledgebase.json")
-
 	if _, err := os.Stat(poolfile); err != nil {
 		// file does not exist, return a new pool
 		return &InMemoryDatabase{
@@ -56,14 +53,26 @@ func NewInMemoryDB(knowledgebase string, store RAGDB) (*InMemoryDatabase, error)
 	if err != nil {
 		return nil, err
 	}
-	return &InMemoryDatabase{
+	db := &InMemoryDatabase{
 		RAGDB:    store,
 		Database: poolData,
 		path:     poolfile,
-	}, nil
+	}
+
+	if err := db.populateRAGDB(); err != nil {
+		return nil, fmt.Errorf("error populating RAGDB: %w", err)
+	}
+
+	return db, nil
 }
 
-func (db *InMemoryDatabase) PopulateRAGDB() error {
+func (db *InMemoryDatabase) Data() []string {
+	db.Lock()
+	defer db.Unlock()
+	return db.Database
+}
+
+func (db *InMemoryDatabase) populateRAGDB() error {
 	for _, d := range db.Database {
 		if d == "" {
 			// skip empty chunks
@@ -139,7 +148,7 @@ func getWebSitemap(url string) (res []string, err error) {
 	return
 }
 
-func WebsiteToKB(website string, chunkSize int, db *InMemoryDatabase) {
+func WebsiteToKB(website string, chunkSize int, db RAGDB) {
 	content, err := getWebSitemap(website)
 	if err != nil {
 		xlog.Info("Error walking sitemap for website", err)
@@ -150,7 +159,7 @@ func WebsiteToKB(website string, chunkSize int, db *InMemoryDatabase) {
 	StringsToKB(db, chunkSize, content...)
 }
 
-func StringsToKB(db *InMemoryDatabase, chunkSize int, content ...string) {
+func StringsToKB(db RAGDB, chunkSize int, content ...string) {
 	for _, c := range content {
 		chunks := splitParagraphIntoChunks(c, chunkSize)
 		xlog.Info("chunks: ", len(chunks))

example/webui/routes.go

@@ -9,7 +9,7 @@ import (
 	"github.com/mudler/local-agent-framework/agent"
 )
 
-func RegisterRoutes(webapp *fiber.App, pool *AgentPool, db *InMemoryDatabase, app *App) {
+func RegisterRoutes(webapp *fiber.App, pool *AgentPool, app *App) {
 
 	webapp.Use("/public", filesystem.New(filesystem.Config{
 		Root:       http.FS(embeddedFiles),
@@ -36,17 +36,20 @@ func RegisterRoutes(webapp *fiber.App, pool *AgentPool, db *InMemoryDatabase, ap
 
 	webapp.Get("/create", func(c *fiber.Ctx) error {
 		return c.Render("views/create", fiber.Map{
-			"Title":      "Hello, World!",
 			"Actions":    AvailableActions,
 			"Connectors": AvailableConnectors,
 		})
 	})
 
-	webapp.Get("/knowledgebase", func(c *fiber.Ctx) error {
+	webapp.Get("/knowledgebase/:name", func(c *fiber.Ctx) error {
-		return c.Render("views/knowledgebase", fiber.Map{
+		db := pool.GetAgentMemory(c.Params("name"))
-			"Title":                   "Hello, World!",
+		return c.Render(
-			"KnowledgebaseItemsCount": len(db.Database),
+			"views/knowledgebase",
-		})
+			fiber.Map{
+				"KnowledgebaseItemsCount": db.Count(),
+				"Name":                    c.Params("name"),
+			},
+		)
 	})
 
 	// Define a route for the GET method on the root path '/'
@@ -80,9 +83,11 @@ func RegisterRoutes(webapp *fiber.App, pool *AgentPool, db *InMemoryDatabase, ap
 	webapp.Put("/pause/:name", app.Pause(pool))
 	webapp.Put("/start/:name", app.Start(pool))
 
-	webapp.Post("/knowledgebase", app.KnowledgeBase(db))
+	webapp.Post("/knowledgebase/:name", app.KnowledgeBase(pool))
-	webapp.Post("/knowledgebase/upload", app.KnowledgeBaseFile(db))
+	webapp.Post("/knowledgebase/:name/upload", app.KnowledgeBaseFile(pool))
-	webapp.Delete("/knowledgebase/reset", app.KnowledgeBaseReset(db))
+	webapp.Delete("/knowledgebase/:name/reset", app.KnowledgeBaseReset(pool))
+	webapp.Post("/knowledgebase/:name/import", app.KnowledgeBaseImport(pool))
+	webapp.Get("/knowledgebase/:name/export", app.KnowledgeBaseExport(pool))
 
 	webapp.Get("/talk/:name", func(c *fiber.Ctx) error {
 		return c.Render("views/chat", fiber.Map{

example/webui/views/index.html

@@ -4,47 +4,7 @@
     <title>Smart Assistant Dashboard</title>
     {{template "views/partials/header"}}
     <style>
-        .container {
+        
-            max-width: 100%;
-            margin: 0 auto;
-            padding: 20px;
-            text-align: center;
-        }
-        .card-link {
-            text-decoration: none; /* Removes underline from links */
-        }
-        .card {
-            background: rgba(255, 255, 255, 0.1);
-            border-radius: 8px;
-            padding: 20px;
-            margin: 20px auto;
-            text-align: left;
-            width: 90%;
-            transition: transform 0.3s ease, box-shadow 0.3s ease; /* Smooth transition for hover effects */
-            display: block; /* Ensures the link fills the card */
-        }
-        .card:hover {
-            transform: translateY(-5px); /* Slight lift effect */
-            box-shadow: 0 10px 20px rgba(0, 0, 0, 0.2); /* Shadow for depth */
-        }
-        .card h2 {
-            font-size: 1.5em; /* Larger and more prominent */
-            font-weight: bold; /* Ensures boldness */
-            color: white; /* Ensures visibility against the card's background */
-            margin-bottom: 0.5em; /* Space below the heading */
-        }
-        .card a,
-        .card p {
-            color: white; /* Ensures text color is consistent */
-        }
-        .card p {
-            font-size: 1em;
-        }
-        .image-container {
-            display: flex;
-            justify-content: center;
-            margin: 20px 0;
-        }
     </style>
 </head>
 <body class="bg-gray-900 p-4 text-white font-sans">
@@ -60,11 +20,10 @@
                 <p>View and manage your list of agents, including detailed profiles and statistics.</p>
             </div>
         </a>
-        <!-- Card for Knowledgebase Management Page -->
+        <a href="/create" class="card-link">
-        <a href="/knowledgebase" class="card-link">
             <div class="card">
-                <h2>Manage Knowledgebase</h2>
+                <h2>Create</h2>
-                <p>Access and update your knowledgebase to improve agent responses and efficiency.</p>
+                <p>Create a new agent.</p>
             </div>
         </a>
     </div>

example/webui/views/knowledgebase.html

@@ -1,7 +1,7 @@
 <!DOCTYPE html>
 <html lang="en">
 <head>
-    <title>KnowledgeBase</title>
+    <title>Knowledgebase for {{.Name}}</title>
     {{template "views/partials/header"}}
 </head>
 <body class="bg-gray-900 p-4 text-white font-sans">
@@ -12,7 +12,7 @@
     <div class="max-w-4xl mx-auto">
         <!-- Add Site Form -->
         <div class="section-box">
-            <form action="/knowledgebase" method="POST">
+            <form action="/knowledgebase/{{.Name}}" method="POST">
                 <h2>Add sites to KB</h2>
                 <label for="url">URL:</label>
                 <input type="text" name="url" id="url" placeholder="Enter URL here">
@@ -24,7 +24,7 @@
 
         <!-- File Upload Form -->
         <div class="section-box">
-            <form id='form' hx-encoding='multipart/form-data' hx-post='/knowledgebase/upload'>
+            <form id='form' hx-encoding='multipart/form-data' hx-post='/knowledgebase/{{.Name}}/upload'>
                 <h2>Upload File</h2>
                 <label for="file">File:</label>
                 <input type='file' name='file' id='file'>
@@ -37,7 +37,22 @@
 
         <!-- Reset Knowledge Base -->
         <div class="section-box">
-            <button hx-swap="none" hx-trigger="click" hx-delete="/knowledgebase/reset">Reset Knowledge Base</button>
+            <button hx-swap="none" hx-trigger="click" hx-delete="/knowledgebase/{{.Name}}/reset">Reset Knowledge Base</button>
+        </div>
+
+        <div class="section-box">
+            <h2>Export</h2>
+            <a href="/knowledgebase/{{.Name}}/export" >Export</a>
+        </div>
+
+        <div class="section-box">
+            <form id='form' hx-encoding='multipart/form-data' hx-post='/knowledgebase/{{.Name}}/import'>
+                <h2>Import</h2>
+                <label for="file">File:</label>
+                <input type='file' name='file' id='file'>
+                <button type="submit">Upload File</button>
+                <progress id='progress' value='0' max='100'></progress>
+            </form>
         </div>
     </div>
     <script>

example/webui/views/partials/header.html

@@ -177,5 +177,47 @@ select::-webkit-scrollbar-thumb {
     transform: rotate(45deg);
 }
 
+.container {
+            max-width: 100%;
+            margin: 0 auto;
+            padding: 20px;
+            text-align: center;
+        }
+        .card-link {
+            text-decoration: none; /* Removes underline from links */
+        }
+        .card {
+            background: rgba(255, 255, 255, 0.1);
+            border-radius: 8px;
+            padding: 20px;
+            margin: 20px auto;
+            text-align: left;
+            width: 90%;
+            transition: transform 0.3s ease, box-shadow 0.3s ease; /* Smooth transition for hover effects */
+            display: block; /* Ensures the link fills the card */
+        }
+        .card:hover {
+            transform: translateY(-5px); /* Slight lift effect */
+            box-shadow: 0 10px 20px rgba(0, 0, 0, 0.2); /* Shadow for depth */
+        }
+        .card h2 {
+            font-size: 1.5em; /* Larger and more prominent */
+            font-weight: bold; /* Ensures boldness */
+            color: white; /* Ensures visibility against the card's background */
+            margin-bottom: 0.5em; /* Space below the heading */
+        }
+        .card a,
+        .card p {
+            color: white; /* Ensures text color is consistent */
+        }
+        .card p {
+            font-size: 1em;
+        }
+        .image-container {
+            display: flex;
+            justify-content: center;
+            margin: 20px 0;
+        }
+
         
 </style>

example/webui/views/partials/menu.html

@@ -14,9 +14,6 @@
                     <a href="/agents" class="px-3 py-2 rounded-md text-sm font-medium text-gray-300 hover:text-white hover:bg-gray-700 focus:outline-none focus:text-white focus:bg-gray-700">
                         <i class="fas fa-users"></i> Agent list
                     </a>
-                    <a href="/knowledgebase" class="px-3 py-2 rounded-md text-sm font-medium text-gray-300 hover:text-white hover:bg-gray-700 focus:outline-none focus:text-white focus:bg-gray-700">
-                        <i class="fas fa-book"></i> Knowledgebase
-                    </a>
                 </div>
             </div>
         </div>

example/webui/views/settings.html

@@ -49,7 +49,12 @@
         <h1 class="text-3xl md:text-5xl font-bold">Agent settings - {{.Name}}</h1>
     </header>
     <div class="max-w-4xl mx-auto">
-
+        <a href="/knowledgebase/{{.Name}}" class="card-link">
+            <div class="card">
+                <h2>Manage Knowledgebase</h2>
+                <p>Access and update your knowledgebase to improve agent responses and efficiency.</p>
+            </div>
+        </a>
         <div class="section-box">
             <button hx-put="/start/{{.Name}}" class="text-indigo-500 hover:text-indigo-400">
                 Start

llm/rag/chromem.go

@@ -42,6 +42,10 @@ func NewChromemDB(collection, path string, openaiClient *openai.Client, embeddin
 	return chromem, nil
 }
 
+func (c *ChromemDB) Count() int {
+	return c.collection.Count()
+}
+
 func (c *ChromemDB) Reset() error {
 	if err := c.db.DeleteCollection(c.collectionName); err != nil {
 		return err

llm/rag/localai.go

@@ -23,6 +23,10 @@ func (db *LocalAIRAGDB) Reset() error {
 	return fmt.Errorf("not implemented")
 }
 
+func (db *LocalAIRAGDB) Count() int {
+	return 0
+}
+
 func (db *LocalAIRAGDB) Store(s string) error {
 	resp, err := db.openaiClient.CreateEmbeddings(context.TODO(),
 		openai.EmbeddingRequestStrings{