diff --git a/Makefile b/Makefile index aa4dac1..ccf7b65 100644 --- a/Makefile +++ b/Makefile @@ -4,6 +4,9 @@ IMAGE_NAME?=webui tests: $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --fail-fast -v -r ./... +webui-nokb: + $(MAKE) webui KBDISABLEINDEX=true + webui: cd example/webui && $(GOCMD) run ./ diff --git a/agent/agent.go b/agent/agent.go index 072ee60..a2d4321 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -38,6 +38,7 @@ type Agent struct { type RAGDB interface { Store(s string) error + Reset() error Search(s string, similarEntries int) ([]string, error) } diff --git a/example/webui/app.go b/example/webui/app.go index cb6ba56..23e036b 100644 --- a/example/webui/app.go +++ b/example/webui/app.go @@ -20,6 +20,13 @@ type ( } ) +func (a *App) KnowledgeBaseReset(db *InMemoryDatabase) func(c *fiber.Ctx) error { + return func(c *fiber.Ctx) error { + db.Reset() + return c.Redirect("/knowledgebase") + } +} + func (a *App) KnowledgeBaseFile(db *InMemoryDatabase) func(c *fiber.Ctx) error { return func(c *fiber.Ctx) error { // https://golang.withcodeexample.com/blog/file-upload-handling-golang-fiber-guide/ diff --git a/example/webui/rag.go b/example/webui/rag.go index 1d90e8f..5fa8b87 100644 --- a/example/webui/rag.go +++ b/example/webui/rag.go @@ -74,6 +74,17 @@ func (db *InMemoryDatabase) SaveToStore() error { return nil } + +func (db *InMemoryDatabase) Reset() error { + db.Lock() + db.Database = []string{} + db.Unlock() + if err := db.rag.Reset(); err != nil { + return err + } + return db.SaveDB() +} + func (db *InMemoryDatabase) AddEntry(entry string) error { db.Lock() defer db.Unlock() diff --git a/example/webui/routes.go b/example/webui/routes.go index 8038ca7..c57684b 100644 --- a/example/webui/routes.go +++ b/example/webui/routes.go @@ -58,8 +58,10 @@ func RegisterRoutes(webapp *fiber.App, pool *AgentPool, db *InMemoryDatabase, ap webapp.Post("/chat/:name", app.Chat(pool)) webapp.Post("/create", app.Create(pool)) webapp.Get("/delete/:name", app.Delete(pool)) + webapp.Post("/knowledgebase", app.KnowledgeBase(db)) webapp.Post("/knowledgebase/upload", app.KnowledgeBaseFile(db)) + webapp.Delete("/knowledgebase/reset", app.KnowledgeBaseReset(db)) webapp.Get("/talk/:name", func(c *fiber.Ctx) error { return c.Render("views/chat", fiber.Map{ diff --git a/example/webui/views/index.html b/example/webui/views/index.html index 6c0b2eb..b03fbbf 100644 --- a/example/webui/views/index.html +++ b/example/webui/views/index.html @@ -3,24 +3,70 @@ Smart Assistant Dashboard {{template "views/partials/header"}} + {{template "views/partials/menu"}} -
-
-

Smart Agent Dashboard

-
-
- - - - Agent List + -
diff --git a/example/webui/views/knowledgebase.html b/example/webui/views/knowledgebase.html index 7019825..1e896c5 100644 --- a/example/webui/views/knowledgebase.html +++ b/example/webui/views/knowledgebase.html @@ -6,55 +6,86 @@ + - + {{template "views/partials/menu"}}

Knowledgebase (items: {{.KnowledgebaseItemsCount}})

-
-
- Add sites to KB -
- - -
-
- - +
+ +
+ +

Add sites to KB

+ + + + + +
- -
- + +
+
+

Upload File

+ + + + + + +
- -
-
- - + +
+
-
- - -
-
- -
-
- -
- +
-
diff --git a/example/webui/views/partials/menu.html b/example/webui/views/partials/menu.html index 7019768..5a964f1 100644 --- a/example/webui/views/partials/menu.html +++ b/example/webui/views/partials/menu.html @@ -2,9 +2,14 @@
+
+ + + Logo + - Home + LocalAgent Agent list diff --git a/llm/rag/chromem.go b/llm/rag/chromem.go index 5f33d2d..617e769 100644 --- a/llm/rag/chromem.go +++ b/llm/rag/chromem.go @@ -13,6 +13,8 @@ type ChromemDB struct { collectionName string collection *chromem.Collection index int + client *openai.Client + db *chromem.DB } func NewChromemDB(collection, path string, openaiClient *openai.Client) (*ChromemDB, error) { @@ -22,9 +24,39 @@ func NewChromemDB(collection, path string, openaiClient *openai.Client) (*Chrome // } db := chromem.NewDB() - embeddingFunc := chromem.EmbeddingFunc( + chromem := &ChromemDB{ + collectionName: collection, + index: 1, + db: db, + client: openaiClient, + } + + c, err := db.GetOrCreateCollection(collection, nil, chromem.embedding()) + if err != nil { + return nil, err + } + chromem.collection = c + + return chromem, nil +} + +func (c *ChromemDB) Reset() error { + if err := c.db.DeleteCollection(c.collectionName); err != nil { + return err + } + collection, err := c.db.GetOrCreateCollection(c.collectionName, nil, c.embedding()) + if err != nil { + return err + } + c.collection = collection + + return nil +} + +func (c *ChromemDB) embedding() chromem.EmbeddingFunc { + return chromem.EmbeddingFunc( func(ctx context.Context, text string) ([]float32, error) { - resp, err := openaiClient.CreateEmbeddings(ctx, + resp, err := c.client.CreateEmbeddings(ctx, openai.EmbeddingRequestStrings{ Input: []string{text}, Model: openai.AdaEmbeddingV2, @@ -43,17 +75,6 @@ func NewChromemDB(collection, path string, openaiClient *openai.Client) (*Chrome return embedding, nil }, ) - - c, err := db.GetOrCreateCollection(collection, nil, embeddingFunc) - if err != nil { - return nil, err - } - - return &ChromemDB{ - collectionName: collection, - collection: c, - index: 1, - }, nil } func (c *ChromemDB) Store(s string) error { diff --git a/llm/rag/localai.go b/llm/rag/localai.go index 7e8bf3a..1e1ecb6 100644 --- a/llm/rag/localai.go +++ b/llm/rag/localai.go @@ -19,6 +19,10 @@ func NewLocalAIRAGDB(storeClient *StoreClient, openaiClient *openai.Client) *Loc } } +func (db *LocalAIRAGDB) Reset() error { + return fmt.Errorf("not implemented") +} + func (db *LocalAIRAGDB) Store(s string) error { resp, err := db.openaiClient.CreateEmbeddings(context.TODO(), openai.EmbeddingRequestStrings{