alex/LocalAGI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support separate knowledge bases for each agent

Browse Source 
Also allow to export/import KB

Signed-off-by: mudler <mudler@localai.io>
This commit is contained in:
mudler
2024-12-18 20:17:05 +01:00
parent c68ff23b01
commit 8c447a0cf8
13 changed files with 205 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
agent/agent.go
View File

@@ -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 {

26
example/webui/agentpool.go
View File

@@ -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 {

70
example/webui/app.go
View File

@@ -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"))
}
}

19
example/webui/main.go
View File

@@ -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"))
}

27
example/webui/rag.go
View File

@@ -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
}
func (db *InMemoryDatabase) PopulateRAGDB() error {
if err := db.populateRAGDB(); err != nil {
return nil, fmt.Errorf("error populating RAGDB: %w", err)
}
return db, nil
}
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))

25
example/webui/routes.go
View File

@@ -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 {
return c.Render("views/knowledgebase", fiber.Map{
"Title": "Hello, World!",
"KnowledgebaseItemsCount": len(db.Database),
})
webapp.Get("/knowledgebase/:name", func(c *fiber.Ctx) error {
db := pool.GetAgentMemory(c.Params("name"))
return c.Render(
"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/upload", app.KnowledgeBaseFile(db))
webapp.Delete("/knowledgebase/reset", app.KnowledgeBaseReset(db))
webapp.Post("/knowledgebase/:name", app.KnowledgeBase(pool))
webapp.Post("/knowledgebase/:name/upload", app.KnowledgeBaseFile(pool))
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{

49
example/webui/views/index.html
View File

@@ -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="/knowledgebase" class="card-link">
<a href="/create" class="card-link">
<div class="card">
<h2>Manage Knowledgebase</h2>
<p>Access and update your knowledgebase to improve agent responses and efficiency.</p>
<h2>Create</h2>
<p>Create a new agent.</p>
</div>
</a>
</div>

23
example/webui/views/knowledgebase.html
View File

@@ -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>

42
example/webui/views/partials/header.html
View File

@@ -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>

3
example/webui/views/partials/menu.html
View File

@@ -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>

7
example/webui/views/settings.html
View File

@@ -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

4
llm/rag/chromem.go
View File

@@ -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

4
llm/rag/localai.go
View File

@@ -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{