From 3891d1c73b404dc0696e9a9450be794510b9782c Mon Sep 17 00:00:00 2001
From: tzwang <wtz1003@sina.com>
Date: Sat, 22 Jun 2024 15:38:17 +0800
Subject: [PATCH] updated imageinfer api

Former-commit-id: 2f5167a75b39f49531c6ebf4e8f32354fcbd44b8
---
 .../logic/inference/imageinferencelogic.go    | 24 +++++++++++++++++++
 pkg/models/taskaisubmodel_gen.go              | 12 +++++-----
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/api/internal/logic/inference/imageinferencelogic.go b/api/internal/logic/inference/imageinferencelogic.go
index 0eac524d..0504be48 100644
--- a/api/internal/logic/inference/imageinferencelogic.go
+++ b/api/internal/logic/inference/imageinferencelogic.go
@@ -292,6 +292,26 @@ func infer(opt *option.InferOption, clusters []*strategy.AssignedCluster, ts []s
 		results = append(results, s)
 	}
 
+	//save ai sub tasks
+	for _, r := range results {
+		for _, task := range aiTaskList {
+			if r.ClusterId == strconv.Itoa(int(task.ClusterId)) {
+				taskAiSub := &models.TaskAiSub{
+					Id:          task.Id,
+					ImageName:   r.ImageName,
+					Result:      r.ImageResult,
+					Card:        r.Card,
+					ClusterId:   task.ClusterId,
+					ClusterName: r.ClusterName,
+				}
+				tx := svcCtx.DbEngin.Save(&taskAiSub)
+				if tx.Error != nil {
+					logx.Errorf(err.Error())
+				}
+			}
+		}
+	}
+
 	sort.Slice(results, func(p, q int) bool {
 		return results[p].ClusterName < results[q].ClusterName
 	})
@@ -345,6 +365,7 @@ func sendInferReq(images []struct {
 				r, err := getInferResult(c.urls[0].Url, t.file, t.imageResult.ImageName, c.clusterName)
 				if err != nil {
 					t.imageResult.ImageResult = err.Error()
+					t.imageResult.ClusterId = c.clusterId
 					t.imageResult.ClusterName = c.clusterName
 					t.imageResult.Card = c.urls[0].Card
 					ch <- t.imageResult
@@ -352,6 +373,7 @@ func sendInferReq(images []struct {
 					return
 				}
 				t.imageResult.ImageResult = r
+				t.imageResult.ClusterId = c.clusterId
 				t.imageResult.ClusterName = c.clusterName
 				t.imageResult.Card = c.urls[0].Card
 
@@ -363,6 +385,7 @@ func sendInferReq(images []struct {
 				r, err := getInferResult(c.urls[idx].Url, t.file, t.imageResult.ImageName, c.clusterName)
 				if err != nil {
 					t.imageResult.ImageResult = err.Error()
+					t.imageResult.ClusterId = c.clusterId
 					t.imageResult.ClusterName = c.clusterName
 					t.imageResult.Card = c.urls[idx].Card
 					ch <- t.imageResult
@@ -370,6 +393,7 @@ func sendInferReq(images []struct {
 					return
 				}
 				t.imageResult.ImageResult = r
+				t.imageResult.ClusterId = c.clusterId
 				t.imageResult.ClusterName = c.clusterName
 				t.imageResult.Card = c.urls[idx].Card
 
diff --git a/pkg/models/taskaisubmodel_gen.go b/pkg/models/taskaisubmodel_gen.go
index 42ff3ed0..3784442e 100644
--- a/pkg/models/taskaisubmodel_gen.go
+++ b/pkg/models/taskaisubmodel_gen.go
@@ -35,12 +35,12 @@ type (
 	}
 
 	TaskAiSub struct {
-		Id          int64          `db:"id"`           // id
-		ImageName   string         `db:"image_name"`   // 图片名称
-		Result      string         `db:"result"`       // 识别结果
-		Card        sql.NullString `db:"card"`         // 加速卡
-		ClusterId   int64          `db:"cluster_id"`   // 集群id
-		ClusterName sql.NullString `db:"cluster_name"` // 集群名称
+		Id          int64  `db:"id"`           // id
+		ImageName   string `db:"image_name"`   // 图片名称
+		Result      string `db:"result"`       // 识别结果
+		Card        string `db:"card"`         // 加速卡
+		ClusterId   int64  `db:"cluster_id"`   // 集群id
+		ClusterName string `db:"cluster_name"` // 集群名称
 	}
 )