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

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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