From 0fa8f6e6e4ba8cca64a8471fd3fc8ca50fa72511 Mon Sep 17 00:00:00 2001
From: Liujian <824010343@qq.com>
Date: Mon, 17 Feb 2025 16:12:21 +0800
Subject: [PATCH] Fix: Service List AI Entry Error

---
 controller/ai-local/iml.go   | 14 ++++----------
 module/ai-local/dto/input.go |  5 +++--
 module/ai-local/iml.go       |  9 +++++++++
 service/ai-local/iml.go      | 12 ++++++++++++
 service/ai-local/service.go  |  1 +
 5 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/controller/ai-local/iml.go b/controller/ai-local/iml.go
index 10a6e860..f6a4c825 100644
--- a/controller/ai-local/iml.go
+++ b/controller/ai-local/iml.go
@@ -83,15 +83,8 @@ func (i *imlLocalModelController) Deploy(ctx *gin.Context) {
 			"code": -1, "msg": "model is required", "success": "fail",
 		})
 		return
-
 	}
-	//err = i.initAILocalService(ctx, input.Model, input.Team)
-	//if err != nil {
-	//	ctx.JSON(200, gin.H{
-	//		"code": -1, "msg": err.Error(), "success": "fail",
-	//	})
-	//	return
-	//}
+
 	id := uuid.NewString()
 	p, err := i.module.Deploy(ctx, input.Model, id)
 	if err != nil {
@@ -105,10 +98,11 @@ func (i *imlLocalModelController) Deploy(ctx *gin.Context) {
 	go func() {
 		select {
 		case <-ctx.Writer.CloseNotify():
-			log.Info("client closed connection,close pipeline")
-			ai_provider_local.CancelPipeline(input.Model, id)
+
 		case <-done:
+
 		}
+		ai_provider_local.CancelPipeline(input.Model, id)
 	}()
 	var complete int64
 	var total int64
diff --git a/module/ai-local/dto/input.go b/module/ai-local/dto/input.go
index b7a3ac98..28ac175f 100644
--- a/module/ai-local/dto/input.go
+++ b/module/ai-local/dto/input.go
@@ -9,6 +9,7 @@ type CancelDeploy struct {
 }
 
 type DeployInput struct {
-	Model string `json:"model"`
-	Team  string `json:"team"`
+	Model   string `json:"model"`
+	Service string `json:"service"`
+	Team    string `json:"team"`
 }
diff --git a/module/ai-local/iml.go b/module/ai-local/iml.go
index b32395be..da22e6b6 100644
--- a/module/ai-local/iml.go
+++ b/module/ai-local/iml.go
@@ -299,6 +299,15 @@ func (i *imlLocalModel) syncGateway(ctx context.Context, clusterId string, relea
 func (i *imlLocalModel) Deploy(ctx context.Context, model string, session string) (*ai_provider_local.Pipeline, error) {
 	var p *ai_provider_local.Pipeline
 	err := i.transaction.Transaction(ctx, func(txCtx context.Context) error {
+		item, err := i.localModelCacheService.GetByTarget(ctx, ai_local.CacheTypeService, model)
+		if err != nil {
+			if !errors.Is(err, gorm.ErrRecordNotFound) {
+				return err
+			}
+
+		} else {
+			model = item.Model
+		}
 		info, err := i.localModelService.Get(ctx, model)
 		if err != nil {
 			if !errors.Is(err, gorm.ErrRecordNotFound) {
diff --git a/service/ai-local/iml.go b/service/ai-local/iml.go
index e037d0a8..7c992c36 100644
--- a/service/ai-local/iml.go
+++ b/service/ai-local/iml.go
@@ -201,6 +201,18 @@ type imlLocalModelCacheService struct {
 	store ai.ILocalModelCacheStore `autowired:""`
 }
 
+func (i *imlLocalModelCacheService) GetByTarget(ctx context.Context, typ CacheType, target string) (*LocalModelCache, error) {
+	item, err := i.store.First(ctx, map[string]interface{}{"target": target, "type": typ.Int()})
+	if err != nil {
+		return nil, err
+	}
+	return &LocalModelCache{
+		Model:  item.Model,
+		Target: item.Target,
+		Type:   CacheType(item.Type),
+	}, nil
+}
+
 func (i *imlLocalModelCacheService) List(ctx context.Context, model string, typ CacheType) ([]*LocalModelCache, error) {
 	list, err := i.store.List(ctx, map[string]interface{}{"model": model, "type": typ.Int()})
 	if err != nil {
diff --git a/service/ai-local/service.go b/service/ai-local/service.go
index 260d3029..e32458ea 100644
--- a/service/ai-local/service.go
+++ b/service/ai-local/service.go
@@ -34,6 +34,7 @@ type ILocalModelCacheService interface {
 	List(ctx context.Context, model string, typ CacheType) ([]*LocalModelCache, error)
 	Delete(ctx context.Context, model string) error
 	Save(ctx context.Context, model string, typ CacheType, target string) error
+	GetByTarget(ctx context.Context, typ CacheType, target string) (*LocalModelCache, error)
 }
 
 func init() {