Skip to content

Commit

Permalink
Wire up a simple explorer DB
Browse files Browse the repository at this point in the history
  • Loading branch information
mudler committed Aug 6, 2024
1 parent 8ece863 commit a404d82
Show file tree
Hide file tree
Showing 5 changed files with 227 additions and 6 deletions.
15 changes: 11 additions & 4 deletions core/cli/explorer.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
107 changes: 107 additions & 0 deletions core/explorer/database.go
Original file line number Diff line number Diff line change
@@ -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)
}
92 changes: 92 additions & 0 deletions core/explorer/database_test.go
Original file line number Diff line number Diff line change
@@ -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)
})
})
})
13 changes: 13 additions & 0 deletions core/explorer/explorer_suite_test.go
Original file line number Diff line number Diff line change
@@ -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")
}
6 changes: 4 additions & 2 deletions core/http/explorer.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit a404d82

Please sign in to comment.