ExternalVendorCode/LocalAI
1
0
Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2024-12-24 06:46:39 +00:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

feat(api): list loaded models in /system (#3661)

Browse Source 
feat(api): list loaded models in /system

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto 2024-09-25 18:00:23 +02:00 committed by GitHub
parent 33b2d38dd0
commit a3d69872e3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 16 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
core/http/endpoints/localai/system.go
View File

@ -17,12 +17,14 @@ func SystemInformations(ml *model.ModelLoader, appConfig *config.ApplicationConf
if err != nil {
return err
}
loadedModels := ml.ListModels()
for b := range appConfig.ExternalGRPCBackends {
availableBackends = append(availableBackends, b)
}
return c.JSON(
schema.SystemInformationResponse{
Backends: availableBackends,
Models: loadedModels,
},
)
}

4
core/schema/localai.go
View File

@ -2,6 +2,7 @@ package schema
import (
"github.com/mudler/LocalAI/core/p2p"
"github.com/mudler/LocalAI/pkg/model"
gopsutil "github.com/shirou/gopsutil/v3/process"
)
@ -72,5 +73,6 @@ type P2PNodesResponse struct {
}
type SystemInformationResponse struct {
Backends []string `json:"backends"`
Backends []string `json:"backends"`
Models []model.Model `json:"loaded_models"`
}

7
pkg/model/initializers.go
View File

@ -311,11 +311,11 @@ func (ml *ModelLoader) grpcModel(backend string, o *Options) func(string, string
log.Debug().Msgf("GRPC Service Started")
client = NewModel(serverAddress)
client = NewModel(modelName, serverAddress)
} else {
log.Debug().Msg("external backend is uri")
// address
client = NewModel(uri)
client = NewModel(modelName, uri)
}
} else {
grpcProcess := backendPath(o.assetDir, backend)
@ -352,7 +352,7 @@ func (ml *ModelLoader) grpcModel(backend string, o *Options) func(string, string
log.Debug().Msgf("GRPC Service Started")
client = NewModel(serverAddress)
client = NewModel(modelName, serverAddress)
}
log.Debug().Msgf("Wait for the service to start up")
@ -419,7 +419,6 @@ func (ml *ModelLoader) BackendLoader(opts ...Option) (client grpc.Backend, err e
err := ml.StopGRPC(allExcept(o.model))
if err != nil {
log.Error().Err(err).Str("keptModel", o.model).Msg("error while shutting down all backends except for the keptModel")
return nil, err
}
}

6
pkg/model/loader.go
View File

@ -105,13 +105,13 @@ FILE:
return models, nil
}
func (ml *ModelLoader) ListModels() []*Model {
func (ml *ModelLoader) ListModels() []Model {
ml.mu.Lock()
defer ml.mu.Unlock()
models := []*Model{}
models := []Model{}
for _, model := range ml.models {
models = append(models, model)
models = append(models, *model)
}
return models

4
pkg/model/loader_test.go
View File

@ -63,7 +63,7 @@ var _ = Describe("ModelLoader", func() {
Context("LoadModel", func() {
It("should load a model and keep it in memory", func() {
mockModel = model.NewModel("test.model")
mockModel = model.NewModel("foo", "test.model")
mockLoader := func(modelName, modelFile string) (*model.Model, error) {
return mockModel, nil
@ -88,7 +88,7 @@ var _ = Describe("ModelLoader", func() {
Context("ShutdownModel", func() {
It("should shutdown a loaded model", func() {
mockModel = model.NewModel("test.model")
mockModel = model.NewModel("foo", "test.model")
mockLoader := func(modelName, modelFile string) (*model.Model, error) {
return mockModel, nil

4
pkg/model/model.go
View File

@ -3,12 +3,14 @@ package model
import grpc "github.com/mudler/LocalAI/pkg/grpc"
type Model struct {
ID string `json:"id"`
address string
client grpc.Backend
}
func NewModel(address string) *Model {
func NewModel(ID, address string) *Model {
return &Model{
ID: ID,
address: address,
}
}