Wire up a simple explorer DB

Author: mudler

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

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)
+}

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)
+		})
+	})
+})

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

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