From a3d69872e35e152f29f7888fa9c56b0a797e9723 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@users.noreply.github.com>
Date: Wed, 25 Sep 2024 18:00:23 +0200
Subject: [PATCH] feat(api): list loaded models in `/system` (#3661)

feat(api): list loaded models in /system

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
---
 core/http/endpoints/localai/system.go | 2 ++
 core/schema/localai.go                | 4 +++-
 pkg/model/initializers.go             | 7 +++----
 pkg/model/loader.go                   | 6 +++---
 pkg/model/loader_test.go              | 4 ++--
 pkg/model/model.go                    | 4 +++-
 6 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/core/http/endpoints/localai/system.go b/core/http/endpoints/localai/system.go
index 11704933..23a725e3 100644
--- 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
index 9070c2be..75fa40c7 100644
--- 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
index 7099bf33..80dd10b4 100644
--- 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
index f70d2cea..4f1ec841 100644
--- 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
index 4621844e..c16a6e50 100644
--- 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
index 1927dc0c..6cb81d10 100644
--- 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,
 	}
 }