From 6cea45e40a7fc9cf3aad5c57d2120d4e40982f25 Mon Sep 17 00:00:00 2001
From: zouap <zouap@pcl.ac.cn>
Date: Fri, 28 Oct 2022 09:17:03 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9F=A5=E8=AF=A2=E6=A8=A1=E5=9E=8B=E6=8E=A5?=
 =?UTF-8?q?=E5=8F=A3=E4=BF=AE=E6=94=B9=E6=88=90=E6=8E=A5=E5=8F=A3=E6=96=87?=
 =?UTF-8?q?=E6=A1=A3=E7=9A=84=E6=96=B9=E5=BC=8F=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: zouap <zouap@pcl.ac.cn>
---
 routers/api/v1/repo/modelmanage.go | 25 ++++++++++++++++++++++++-
 routers/repo/ai_model_manage.go    | 12 +++++++-----
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/routers/api/v1/repo/modelmanage.go b/routers/api/v1/repo/modelmanage.go
index b32726997..b404558a1 100644
--- a/routers/api/v1/repo/modelmanage.go
+++ b/routers/api/v1/repo/modelmanage.go
@@ -8,6 +8,15 @@ import (
 	routerRepo "code.gitea.io/gitea/routers/repo"
 )
 
+type FileInfo struct {
+	FileName string `json:"fileName"`
+	ModTime  string `json:"modTime"`
+	IsDir    bool   `json:"isDir"`
+	Size     int64  `json:"size"`
+	ParenDir string `json:"parenDir"`
+	UUID     string `json:"uuid"`
+}
+
 func CreateNewModel(ctx *context.APIContext) {
 	log.Info("CreateNewModel by api.")
 	routerRepo.SaveModel(ctx.Context)
@@ -35,7 +44,21 @@ func QueryModelListForPredict(ctx *context.APIContext) {
 
 func QueryModelFileForPredict(ctx *context.APIContext) {
 	log.Info("QueryModelFileForPredict by api.")
-	routerRepo.QueryModelFileForPredict(ctx.Context)
+	id := ctx.Query("id")
+	result := routerRepo.QueryModelFileByID(id)
+	re := make([]FileInfo, len(result))
+	for _, file := range result {
+		tmpFile := FileInfo{
+			FileName: file.FileName,
+			ModTime:  file.ModTime,
+			IsDir:    file.IsDir,
+			Size:     file.Size,
+			ParenDir: file.ParenDir,
+			UUID:     file.UUID,
+		}
+		re = append(re, tmpFile)
+	}
+	ctx.JSON(http.StatusOK, re)
 }
 
 func CreateModelConvert(ctx *context.APIContext) {
diff --git a/routers/repo/ai_model_manage.go b/routers/repo/ai_model_manage.go
index 6385cc674..45c21d719 100644
--- a/routers/repo/ai_model_manage.go
+++ b/routers/repo/ai_model_manage.go
@@ -910,23 +910,25 @@ func QueryModelFileForPredict(ctx *context.Context) {
 	if id == "" {
 		id = ctx.Query("ID")
 	}
+	ctx.JSON(http.StatusOK, QueryModelFileByID(id))
+}
+
+func QueryModelFileByID(id string) []storage.FileInfo {
 	model, err := models.QueryModelById(id)
 	if err == nil {
 		if model.Type == models.TypeCloudBrainTwo {
 			prefix := model.Path[len(setting.Bucket)+1:]
 			fileinfos, _ := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, prefix)
-			ctx.JSON(http.StatusOK, fileinfos)
+			return fileinfos
 		} else if model.Type == models.TypeCloudBrainOne {
 			prefix := model.Path[len(setting.Attachment.Minio.Bucket)+1:]
 			fileinfos, _ := storage.GetAllObjectByBucketAndPrefixMinio(setting.Attachment.Minio.Bucket, prefix)
-			ctx.JSON(http.StatusOK, fileinfos)
+			return fileinfos
 		}
 	} else {
 		log.Error("no such model!", err.Error())
-		ctx.ServerError("no such model:", err)
-		return
 	}
-
+	return nil
 }
 
 func QueryOneLevelModelFile(ctx *context.Context) {