Jimeng VirtualHuman and actionTransfer is ready

api/service/jimeng/client.go

@@ -3,11 +3,13 @@ package jimeng
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"geekai/core/types"
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 
 	"github.com/volcengine/volc-sdk-golang/base"
 	"github.com/volcengine/volc-sdk-golang/service/visual"
@@ -54,6 +56,22 @@ func (c *Client) UpdateConfig(config types.JimengConfig) error {
 				"Version": []string{"2022-08-31"},
 			},
 		},
+		"CVSubmitTask": {
+			Method: http.MethodPost,
+			Path:   "/",
+			Query: url.Values{
+				"Action":  []string{"CVSubmitTask"},
+				"Version": []string{"2022-08-31"},
+			},
+		},
+		"CVGetResult": {
+			Method: http.MethodPost,
+			Path:   "/",
+			Query: url.Values{
+				"Action":  []string{"CVGetResult"},
+				"Version": []string{"2022-08-31"},
+			},
+		},
 		"CVProcess": {
 			Method: http.MethodPost,
 			Path:   "/",
@@ -75,6 +93,22 @@ func (c *Client) UpdateConfig(config types.JimengConfig) error {
 	return c.testConnection()
 }
 
+// GetErrorMessage 根据错误代码获取对应的错误信息
+func GetErrorMessage(code int) string {
+	if message, exists := errorCodeMessages[code]; exists {
+		return message
+	}
+	return fmt.Sprintf("未知错误代码: %d", code)
+}
+
+// HandleResponseError 处理响应错误，根据错误代码返回详细的错误信息
+func HandleResponseError(code int, message string) error {
+	if code == ECSuccess {
+		return nil
+	}
+	return errors.New(GetErrorMessage(code))
+}
+
 // testConnection 测试即梦AI连接
 func (c *Client) testConnection() error {
 
@@ -84,7 +118,7 @@ func (c *Client) testConnection() error {
 		TaskId: "test_task_id_12345",
 	}
 
-	_, err := c.QueryTask(testReq)
+	_, err := c.QueryTask(testReq, ASyncActionGetResult)
 	// 即使任务不存在，只要不是认证错误就说明连接正常
 	if err != nil {
 		// 检查是否是认证错误
@@ -105,17 +139,16 @@ func (c *Client) SubmitTask(req map[string]any) (*SubmitTaskResponse, error) {
 		return nil, fmt.Errorf("marshal request failed: %w", err)
 	}
 
-	// 单独处理图片特效任务
-	if req["req_key"] == ImageEffectReqKey {
-		req["image_input1"] = req["image_urls"].([]any)[0]
-		delete(req, "image_urls")
-	}
-
 	// 直接使用序列化后的字节
 	jsonBody := reqBodyBytes
+	action := ASyncActionSubmit
+	if v, ok := req["action"]; ok {
+		action = v.(string)
+		delete(req, "action")
+	}
 
 	// 调用SDK的JSON方法
-	respBody, statusCode, err := c.visual.Client.Json("CVSync2AsyncSubmitTask", nil, string(jsonBody))
+	respBody, statusCode, err := c.visual.Client.Json(action, nil, string(jsonBody))
 	if err != nil {
 		return nil, fmt.Errorf("submit task failed (status: %d): %w", statusCode, err)
 	}
@@ -128,11 +161,70 @@ func (c *Client) SubmitTask(req map[string]any) (*SubmitTaskResponse, error) {
 		return nil, fmt.Errorf("unmarshal response failed: %w", err)
 	}
 
+	// 检查响应错误代码
+	if err := HandleResponseError(result.Code, result.Message); err != nil {
+		return nil, err
+	}
+
 	return &result, nil
 }
 
+// 识别数字人主体
+func (c *Client) AvatarRecognition(imgUrl string, reqKey string) error {
+	params := map[string]any{
+		"image_url": imgUrl,
+		"req_key":   reqKey,
+	}
+	reqBodyBytes, err := json.Marshal(params)
+	if err != nil {
+		return fmt.Errorf("marshal request failed: %w", err)
+	}
+	// 调用SDK的JSON方法
+	respBody, statusCode, err := c.visual.Client.Json(SyncActionSubmit, nil, string(reqBodyBytes))
+	if err != nil {
+		return fmt.Errorf("submit task failed (status: %d): %w", statusCode, err)
+	}
+
+	// 解析响应
+	var result SubmitTaskResponse
+	if err := json.Unmarshal(respBody, &result); err != nil {
+		return fmt.Errorf("unmarshal response failed: %w", err)
+	}
+
+	// 检查响应错误代码
+	if err := HandleResponseError(result.Code, result.Message); err != nil {
+		return err
+	}
+
+	// 等待任务完成
+	for {
+		resp, err := c.QueryTask(&QueryTaskRequest{
+			ReqKey: reqKey,
+			TaskId: result.Data.TaskId,
+		}, SyncActionGetResult)
+		if err != nil {
+			return fmt.Errorf("query task failed: %w", err)
+		}
+		if resp.Data.Status != types.JMTaskStatusDone {
+			time.Sleep(time.Second * 3)
+			continue
+		}
+		var respData map[string]int
+		if err := json.Unmarshal([]byte(resp.Data.RespData), &respData); err != nil {
+			return fmt.Errorf("unmarshal response failed: %w", err)
+		}
+		logger.Debugf("Jimeng AvatarRecognition Response: %+v", resp)
+		if respData["status"] == 1 {
+			return nil
+		} else {
+			return errors.New("不包含人、类人、拟人等主体")
+		}
+
+	}
+}
+
 // QueryTask 查询任务结果
-func (c *Client) QueryTask(req *QueryTaskRequest) (*QueryTaskResponse, error) {
+func (c *Client) QueryTask(req *QueryTaskRequest, action string) (*QueryTaskResponse, error) {
 	// 序列化请求
 	jsonBody, err := json.Marshal(req)
 	if err != nil {
@@ -140,7 +232,7 @@ func (c *Client) QueryTask(req *QueryTaskRequest) (*QueryTaskResponse, error) {
 	}
 
 	// 调用SDK的JSON方法
-	respBody, statusCode, err := c.visual.Client.Json("CVSync2AsyncGetResult", nil, string(jsonBody))
+	respBody, statusCode, err := c.visual.Client.Json(action, nil, string(jsonBody))
 	if err != nil {
 		return nil, fmt.Errorf("query task failed (status: %d): %w", statusCode, err)
 	}
@@ -153,6 +245,11 @@ func (c *Client) QueryTask(req *QueryTaskRequest) (*QueryTaskResponse, error) {
 		return nil, fmt.Errorf("unmarshal response failed: %w", err)
 	}
 
+	// 检查响应错误代码
+	if err := HandleResponseError(result.Code, result.Message); err != nil {
+		return nil, err
+	}
+
 	return &result, nil
 }