From a404d822878261b7294532ed52bca1b6443669b3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 6 Aug 2024 12:08:17 +0200 Subject: [PATCH] Wire up a simple explorer DB --- core/cli/explorer.go | 15 +++- core/explorer/database.go | 107 +++++++++++++++++++++++++++ core/explorer/database_test.go | 92 +++++++++++++++++++++++ core/explorer/explorer_suite_test.go | 13 ++++ core/http/explorer.go | 6 +- 5 files changed, 227 insertions(+), 6 deletions(-) create mode 100644 core/explorer/database.go create mode 100644 core/explorer/database_test.go create mode 100644 core/explorer/explorer_suite_test.go diff --git a/core/cli/explorer.go b/core/cli/explorer.go index e0f92be81dec..9e71c508feb3 100644 --- a/core/cli/explorer.go +++ b/core/cli/explorer.go @@ -2,15 +2,22 @@ package cli import ( cliContext "github.com/mudler/LocalAI/core/cli/context" + "github.com/mudler/LocalAI/core/explorer" "github.com/mudler/LocalAI/core/http" ) type ExplorerCMD struct { - Address string `env:"LOCALAI_ADDRESS,ADDRESS" default:":8080" help:"Bind address for the API server" group:"api"` + Address string `env:"LOCALAI_ADDRESS,ADDRESS" default:":8080" help:"Bind address for the API server" group:"api"` + PoolDatabase string `env:"LOCALAI_POOL_DATABASE,POOL_DATABASE" default:"" help:"Path to the pool database" group:"api"` } -func (explorer *ExplorerCMD) Run(ctx *cliContext.Context) error { - appHTTP := http.Explorer() +func (e *ExplorerCMD) Run(ctx *cliContext.Context) error { - return appHTTP.Listen(explorer.Address) + db, err := explorer.NewDatabase(e.PoolDatabase) + if err != nil { + return err + } + appHTTP := http.Explorer(db) + + return appHTTP.Listen(e.Address) } diff --git a/core/explorer/database.go b/core/explorer/database.go new file mode 100644 index 000000000000..092623349c03 --- /dev/null +++ b/core/explorer/database.go @@ -0,0 +1,107 @@ +package explorer + +// A simple JSON database for storing and retrieving p2p network tokens and a name and description. + +import ( + "encoding/json" + "os" + "sort" + "sync" +) + +// Database is a simple JSON database for storing and retrieving p2p network tokens and a name and description. +type Database struct { + sync.RWMutex + path string + data map[string]TokenData +} + +// TokenData is a p2p network token with a name and description. +type TokenData struct { + Name string `json:"name"` + Description string `json:"description"` +} + +// NewDatabase creates a new Database with the given path. +func NewDatabase(path string) (*Database, error) { + db := &Database{ + data: make(map[string]TokenData), + path: path, + } + return db, db.load() +} + +// Get retrieves a Token from the Database by its token. +func (db *Database) Get(token string) (TokenData, bool) { + db.RLock() + defer db.RUnlock() + t, ok := db.data[token] + return t, ok +} + +// Set stores a Token in the Database by its token. +func (db *Database) Set(token string, t TokenData) error { + db.Lock() + db.data[token] = t + db.Unlock() + + return db.save() +} + +// Delete removes a Token from the Database by its token. +func (db *Database) Delete(token string) error { + db.Lock() + delete(db.data, token) + db.Unlock() + return db.save() +} + +func (db *Database) TokenList() []string { + db.RLock() + defer db.RUnlock() + tokens := []string{} + for k := range db.data { + tokens = append(tokens, k) + } + + sort.Slice(tokens, func(i, j int) bool { + // sort by token + return tokens[i] < tokens[j] + }) + + return tokens +} + +// load reads the Database from disk. +func (db *Database) load() error { + db.Lock() + defer db.Unlock() + + if _, err := os.Stat(db.path); os.IsNotExist(err) { + + return nil + } + + // Read the file from disk + // Unmarshal the JSON into db.data + f, err := os.ReadFile(db.path) + if err != nil { + return err + } + return json.Unmarshal(f, &db.data) +} + +// save writes the Database to disk. +func (db *Database) save() error { + db.RLock() + defer db.RUnlock() + + // Marshal db.data into JSON + // Write the JSON to the file + f, err := os.Create(db.path) + if err != nil { + return err + } + defer f.Close() + return json.NewEncoder(f).Encode(db.data) +} diff --git a/core/explorer/database_test.go b/core/explorer/database_test.go new file mode 100644 index 000000000000..7f2cbd268a36 --- /dev/null +++ b/core/explorer/database_test.go @@ -0,0 +1,92 @@ +package explorer_test + +import ( + "os" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/explorer" +) + +var _ = Describe("Database", func() { + var ( + dbPath string + db *explorer.Database + err error + ) + + BeforeEach(func() { + // Create a temporary file path for the database + dbPath = "test_db.json" + db, err = explorer.NewDatabase(dbPath) + Expect(err).To(BeNil()) + }) + + AfterEach(func() { + // Clean up the temporary database file + os.Remove(dbPath) + }) + + Context("when managing tokens", func() { + It("should add and retrieve a token", func() { + token := "token123" + t := explorer.TokenData{Name: "TokenName", Description: "A test token"} + + err = db.Set(token, t) + Expect(err).To(BeNil()) + + retrievedToken, exists := db.Get(token) + Expect(exists).To(BeTrue()) + Expect(retrievedToken).To(Equal(t)) + }) + + It("should delete a token", func() { + token := "token123" + t := explorer.TokenData{Name: "TokenName", Description: "A test token"} + + err = db.Set(token, t) + Expect(err).To(BeNil()) + + err = db.Delete(token) + Expect(err).To(BeNil()) + + _, exists := db.Get(token) + Expect(exists).To(BeFalse()) + }) + + It("should persist data to disk", func() { + token := "token123" + t := explorer.TokenData{Name: "TokenName", Description: "A test token"} + + err = db.Set(token, t) + Expect(err).To(BeNil()) + + // Recreate the database object to simulate reloading from disk + db, err = explorer.NewDatabase(dbPath) + Expect(err).To(BeNil()) + + retrievedToken, exists := db.Get(token) + Expect(exists).To(BeTrue()) + Expect(retrievedToken).To(Equal(t)) + + // Check the token list + tokenList := db.TokenList() + Expect(tokenList).To(ContainElement(token)) + }) + }) + + Context("when loading an empty or non-existent file", func() { + It("should start with an empty database", func() { + dbPath = "empty_db.json" + db, err = explorer.NewDatabase(dbPath) + Expect(err).To(BeNil()) + + _, exists := db.Get("nonexistent") + Expect(exists).To(BeFalse()) + + // Clean up + os.Remove(dbPath) + }) + }) +}) diff --git a/core/explorer/explorer_suite_test.go b/core/explorer/explorer_suite_test.go new file mode 100644 index 000000000000..fc718d5f8dfa --- /dev/null +++ b/core/explorer/explorer_suite_test.go @@ -0,0 +1,13 @@ +package explorer_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestExplorer(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Explorer test suite") +} diff --git a/core/http/explorer.go b/core/http/explorer.go index d7f897635075..d69d0934f9d5 100644 --- a/core/http/explorer.go +++ b/core/http/explorer.go @@ -2,20 +2,22 @@ package http import ( "github.com/gofiber/fiber/v2" + "github.com/mudler/LocalAI/core/explorer" "github.com/mudler/LocalAI/core/http/routes" ) -func Explorer() *fiber.App { +func Explorer(db *explorer.Database) *fiber.App { fiberCfg := fiber.Config{ Views: renderEngine(), // We disable the Fiber startup message as it does not conform to structured logging. // We register a startup log line with connection information in the OnListen hook to keep things user friendly though - DisableStartupMessage: true, + DisableStartupMessage: false, // Override default error handler } app := fiber.New(fiberCfg) + routes.RegisterExplorerRoutes(app) return app