feat(api): list loaded models in /system

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
mudler · Sep 25, 2024 · 4a4c975 · 4a4c975
1 parent 33b2d38
commit 4a4c975
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 11 deletions.
diff --git a/core/http/endpoints/localai/system.go b/core/http/endpoints/localai/system.go
@@ -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,
 			},
 		)
 	}

diff --git a/core/schema/localai.go b/core/schema/localai.go
@@ -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"`
 }
diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go
@@ -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
 		}
 	}
 

diff --git a/pkg/model/loader.go b/pkg/model/loader.go
@@ -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

diff --git a/pkg/model/loader_test.go b/pkg/model/loader_test.go
@@ -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

diff --git a/pkg/model/model.go b/pkg/model/model.go
@@ -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,
 	}
 }