parse container resp data

nm/nm_operate.go

@@ -73,6 +73,7 @@ func inputNodeManagerChan(manager *NodeManager, nodeManagerClient *models.NodeMa
 			IsSelected:    isSelect,
 			LastHeartTime: time.Now(),
 		}
+		usedNodeManagerClient = append(usedNodeManagerClient, nodeManagerClient)
 	}
 	serviceClient := operate.ConnNmGrpc(manager.Info.Endpoint)
 	if serviceClient == nil {
@@ -82,7 +83,6 @@ func inputNodeManagerChan(manager *NodeManager, nodeManagerClient *models.NodeMa
 	nodeManagerClient.Status = true
 	nodeManagerClient.Client = serviceClient
 	nodeManagerClientChan <- nodeManagerClient
-	usedNodeManagerClient = append(usedNodeManagerClient, nodeManagerClient)
 	manager.IsUsed = true
 	return true
 }

nm/task_msg.go

@@ -2,7 +2,6 @@ package nm
 
 import (
 	"bytes"
-	"encoding/base64"
 	"encoding/json"
 	"example.com/m/conf"
 	"example.com/m/log"
@@ -233,31 +232,69 @@ func (t *TaskHandler) ComputeTaskHandler(taskMsg *nodeManagerV1.PushTaskMessage)
 			}
 			log.WithField("isUseFileCache", isUseFileCache).Info("is use file cache")
 			if isUseFileCache && readBody != nil {
-				containerResp := &models.ModelResponse{}
-				//log.WithField("task resp", readBody).Info("received response")
-				err = json.Unmarshal(readBody, &containerResp)
-				if err != nil {
-					log.WithError(err).Error("Error unmarshalling oss resp body failed")
-					return
-				}
-				if len(containerResp.Output) == 1 {
-					if utils.IsBase64ImageStr(containerResp.Output[0]) {
-						containerRespOutput := strings.SplitN(containerResp.Output[0], ",", 2)
-						imageStr := containerRespOutput[1]
-						queryString := utils.MatchFileCacheQueryString(taskParam.Headers, taskCmd.ImageName, t.DockerOp.ModelsInfo, containerRespOutput[0])
-						//ossUri, err := t.uploadOSS(taskMsg.TaskId, queryString, containerResp.Output[0])
-						ossUri, err := t.uploadOSS(taskMsg.TaskId, queryString, imageStr)
-						if err != nil {
-							log.WithError(err).Error("upload image into file cache failed")
-							return
+				respBody := make([]byte, 0)
+				base64ImageStr := ""
+				data := parseData(readBody)
+				if data != nil {
+					switch v := data.(type) {
+					case [][]string:
+						{
+							res := data.([][]string)
+							if len(res) == 1 {
+								if len(res[0]) == 1 {
+									base64ImageStr = res[0][0]
+								}
+							}
+							for _, innerSlice := range res {
+								for _, str := range innerSlice {
+									respBody = append(respBody, []byte(str)...)
+								}
+							}
+							log.Info("data is [][]string type")
+						}
+					case []string:
+						{
+							res := data.([]string)
+							if len(res) == 1 {
+								base64ImageStr = res[0]
+							}
+							respBody = []byte(strings.Join(res, ""))
+							log.Info("data is []string type")
 						}
-						log.WithField("uri", ossUri).Info("upload image OSS  successful")
-						if ossUri != "" {
-							taskExecResult.TaskHttpStatusCode = models.RedirectCode
-							post.Header.Set("Location", ossUri)
+					case string:
+						{
+							res := data.(string)
+							base64ImageStr = res
+							respBody = []byte(res)
+							log.Info("data is string type")
 						}
+					default:
+						log.Error("data is unknown type", v)
+					}
+					if base64ImageStr != "" {
+						isBase64, decodeByte, respFormat, suffix := utils.IsBase64ImageStr(base64ImageStr)
+						if isBase64 {
+							log.WithField("taskId", taskMsg.TaskId).WithField("format", respFormat).WithField("suffix", suffix).Info("Parse container resp")
+							queryString := utils.MatchFileCacheQueryString(taskParam.Headers, taskCmd.ImageName, t.DockerOp.ModelsInfo, respFormat)
+							ossUri, err := t.uploadOSS(taskMsg.TaskId, queryString, decodeByte, suffix)
+							if err != nil {
+								log.WithError(err).Error("upload image into file cache failed")
+								return
+							}
+							log.WithField("uri", ossUri).Info("upload image OSS  successful")
+							if ossUri != "" {
+								taskExecResult.TaskHttpStatusCode = models.RedirectCode
+								post.Header.Set("Location", ossUri)
+							}
+						} else {
+							taskExecResult.TaskRespBody = respBody
+						}
+					} else {
+						taskExecResult.TaskRespBody = respBody
 					}
 				}
+			} else {
+				taskExecResult.TaskRespBody = readBody
 			}
 		}
 		headers, err := json.Marshal(post.Header)
@@ -336,23 +373,15 @@ func (t *TaskHandler) foundImageIsRunning(imageId string) (bool, string, uint16)
 	return false, "", 0
 }
 
-func (t *TaskHandler) uploadOSS(taskId string, queries string, base64Image string) (string, error) {
-	// todo: 解析结果
-	// TODO: 存储OSS
+func (t *TaskHandler) uploadOSS(taskId string, queries string, decodedImage []byte, suffix string) (string, error) {
 	var requestBody bytes.Buffer
 	writer := multipart.NewWriter(&requestBody)
 	// 创建文件表单字段
-	fileField, err := writer.CreateFormFile("file", fmt.Sprintf("%s.png", taskId))
+	fileField, err := writer.CreateFormFile("file", fmt.Sprintf("%s.%s", taskId, suffix))
 	if err != nil {
 		log.WithError(err).Error("Error creating form file")
 		return "", err
 	}
-	// 将 base64 解码后的内容复制到表单字段
-	decodedImage, err := base64.StdEncoding.DecodeString(base64Image)
-	if err != nil {
-		log.WithError(err).Error("Error decoding base64 image")
-		return "", err
-	}
 	_, err = io.Copy(fileField, bytes.NewReader(decodedImage))
 	//_, err = io.Copy(fileField, strings.NewReader(base64Image))
 	if err != nil {
@@ -383,3 +412,29 @@ func (t *TaskHandler) uploadOSS(taskId string, queries string, base64Image strin
 	}
 	return bytes.NewBuffer(ossRespBody).String(), nil
 }
+
+func parseData(readBody []byte) interface{} {
+	var m map[string]json.RawMessage
+	if err := json.Unmarshal(readBody, &m); err != nil {
+		log.WithError(err).Error("Parse json raw message failed")
+		return bytes.NewBuffer(readBody).String()
+	}
+	var outputTwoArray [][]string
+	if err := json.Unmarshal(m["output"], &outputTwoArray); err != nil {
+		log.WithField("err", err).Warn("parse two array output filed failed:")
+		var outputOneArray []string
+		if err := json.Unmarshal(m["output"], &outputOneArray); err != nil {
+			log.WithField("err", err).Warn("parse one array output filed failed:")
+			var outputString string
+			if err := json.Unmarshal(m["output"], &outputString); err != nil {
+				return nil
+			} else {
+				return outputString
+			}
+		} else {
+			return outputOneArray
+		}
+	} else {
+		return outputTwoArray
+	}
+}

test/rootcmd_test.go

 package test
 
 import (
+	"encoding/json"
 	"fmt"
-	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/crypto"
 	"testing"
 )
 
 func Test_initConfig(t *testing.T) {
-	prvKey, _ := crypto.HexToECDSA("e3b097b0c171e2489973a277b1546392db97e359505cd64b9b52966cb87a0f08")
-	fmt.Println("prvKey:", prvKey)
-	pubKey := common.Bytes2Hex(crypto.FromECDSAPub(&prvKey.PublicKey))
-	fmt.Println("pubKey:", pubKey)
-	address := crypto.PubkeyToAddress(prvKey.PublicKey)
-	fmt.Println("address:", address)
+	//prvKey, _ := crypto.HexToECDSA("e3b097b0c171e2489973a277b1546392db97e359505cd64b9b52966cb87a0f08")
+	//fmt.Println("prvKey:", prvKey)
+	//pubKey := common.Bytes2Hex(crypto.FromECDSAPub(&prvKey.PublicKey))
+	//fmt.Println("pubKey:", pubKey)
+	//address := crypto.PubkeyToAddress(prvKey.PublicKey)
+	//fmt.Println("address:", address)
+
+	// JSON 2 数据
+	jsonData := `{
+		"completed_at": "2023-07-02T02:13:48.764861Z",
+		"output": ["ss","sss"]
+	}`
+
+	// 解析 JSON 数据到 map[string]json.RawMessage
+	var m map[string]json.RawMessage
+	if err := json.Unmarshal([]byte(jsonData), &m); err != nil {
+		fmt.Println("解析 JSON 数据时出错:", err)
+		return
+	}
+
+	// 解析 "output" 字段
+	//var output [][]string
+	//if err := json.Unmarshal(m["output"], &output); err != nil {
+	//	fmt.Println("解析 output 字段时出错:", err)
+	//	return
+	//}
+
+	var output []string
+	if err := json.Unmarshal(m["output"], &output); err != nil {
+		fmt.Println("解析 output 字段时出错:", err)
+		return
+	}
+
+	// 输出结果
+	fmt.Println("Output Type:", output)
 }

utils/util.go

@@ -96,14 +96,28 @@ func readPrivateKey() (*ecdsa.PrivateKey, error) {
 }
 
 // IsBase64ImageStr 检查字符串是否是 Base64 编码的图像数据
-func IsBase64ImageStr(imageStr string) bool {
+func IsBase64ImageStr(imageStr string) (bool, []byte, string, string) {
 	// 移除可能的前缀（如 "data:image/png;base64,"）
-	imageStr = strings.SplitN(imageStr, ",", 2)[1]
-	_, err := base64.StdEncoding.DecodeString(imageStr)
-	return err == nil
+	imageStrArr := strings.SplitN(imageStr, ",", 2)
+	base64CodePrefix := ""
+	base64Code := ""
+	if len(imageStrArr) == 1 {
+		base64Code = imageStrArr[0]
+		base64CodePrefix = "data:image/png;base64"
+	} else {
+		base64CodePrefix = imageStrArr[0]
+		base64Code = imageStrArr[1]
+	}
+	decodeBytes, err := base64.StdEncoding.DecodeString(base64Code)
+	if err != nil {
+		return false, nil, "", ""
+	}
+	formatStr := strings.Split(strings.Split(base64CodePrefix, ";")[0], ":")[1]
+	suffix := strings.Split(formatStr, "/")[1]
+	return true, decodeBytes, formatStr, suffix
 }
 
-func MatchFileCacheQueryString(params map[string][]string, taskImageName string, modelsInfo []*models.ModelInfo, formatType string) string {
+func MatchFileCacheQueryString(params map[string][]string, taskImageName string, modelsInfo []*models.ModelInfo, contentType string) string {
 	values := url.Values{}
 	isExistFileExpires := false
 	for key, value := range params {
@@ -126,7 +140,6 @@ func MatchFileCacheQueryString(params map[string][]string, taskImageName string,
 	if !isModelExistFileExpires {
 		values.Add(models.ResultFileExpiresDB, "600")
 	}
-	contentType := strings.Split(strings.Split(formatType, ";")[0], ":")[1]
 	values.Add(models.ContentType, contentType)
 	return values.Encode()
 }