chore: fix ali image implementation

relay/channel/ali/image.go

@@ -0,0 +1,192 @@
+package ali
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/songquanpeng/one-api/common/helper"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/relay/channel/openai"
+	"github.com/songquanpeng/one-api/relay/model"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+)
+
+func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) {
+	apiKey := c.Request.Header.Get("Authorization")
+	apiKey = strings.TrimPrefix(apiKey, "Bearer ")
+	responseFormat := c.GetString("response_format")
+
+	var aliTaskResponse TaskResponse
+	responseBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return openai.ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = resp.Body.Close()
+	if err != nil {
+		return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = json.Unmarshal(responseBody, &aliTaskResponse)
+	if err != nil {
+		return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
+	}
+
+	if aliTaskResponse.Message != "" {
+		logger.SysError("aliAsyncTask err: " + string(responseBody))
+		return openai.ErrorWrapper(errors.New(aliTaskResponse.Message), "ali_async_task_failed", http.StatusInternalServerError), nil
+	}
+
+	aliResponse, _, err := asyncTaskWait(aliTaskResponse.Output.TaskId, apiKey)
+	if err != nil {
+		return openai.ErrorWrapper(err, "ali_async_task_wait_failed", http.StatusInternalServerError), nil
+	}
+
+	if aliResponse.Output.TaskStatus != "SUCCEEDED" {
+		return &model.ErrorWithStatusCode{
+			Error: model.Error{
+				Message: aliResponse.Output.Message,
+				Type:    "ali_error",
+				Param:   "",
+				Code:    aliResponse.Output.Code,
+			},
+			StatusCode: resp.StatusCode,
+		}, nil
+	}
+
+	fullTextResponse := responseAli2OpenAIImage(aliResponse, responseFormat)
+	jsonResponse, err := json.Marshal(fullTextResponse)
+	if err != nil {
+		return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil
+	}
+	c.Writer.Header().Set("Content-Type", "application/json")
+	c.Writer.WriteHeader(resp.StatusCode)
+	_, err = c.Writer.Write(jsonResponse)
+	return nil, nil
+}
+
+func asyncTask(taskID string, key string) (*TaskResponse, error, []byte) {
+	url := fmt.Sprintf("https://dashscope.aliyuncs.com/api/v1/tasks/%s", taskID)
+
+	var aliResponse TaskResponse
+
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		return &aliResponse, err, nil
+	}
+
+	req.Header.Set("Authorization", "Bearer "+key)
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		logger.SysError("aliAsyncTask client.Do err: " + err.Error())
+		return &aliResponse, err, nil
+	}
+	defer resp.Body.Close()
+
+	responseBody, err := io.ReadAll(resp.Body)
+
+	var response TaskResponse
+	err = json.Unmarshal(responseBody, &response)
+	if err != nil {
+		logger.SysError("aliAsyncTask NewDecoder err: " + err.Error())
+		return &aliResponse, err, nil
+	}
+
+	return &response, nil, responseBody
+}
+
+func asyncTaskWait(taskID string, key string) (*TaskResponse, []byte, error) {
+	waitSeconds := 2
+	step := 0
+	maxStep := 20
+
+	var taskResponse TaskResponse
+	var responseBody []byte
+
+	for {
+		step++
+		rsp, err, body := asyncTask(taskID, key)
+		responseBody = body
+		if err != nil {
+			return &taskResponse, responseBody, err
+		}
+
+		if rsp.Output.TaskStatus == "" {
+			return &taskResponse, responseBody, nil
+		}
+
+		switch rsp.Output.TaskStatus {
+		case "FAILED":
+			fallthrough
+		case "CANCELED":
+			fallthrough
+		case "SUCCEEDED":
+			fallthrough
+		case "UNKNOWN":
+			return rsp, responseBody, nil
+		}
+		if step >= maxStep {
+			break
+		}
+		time.Sleep(time.Duration(waitSeconds) * time.Second)
+	}
+
+	return nil, nil, fmt.Errorf("aliAsyncTaskWait timeout")
+}
+
+func responseAli2OpenAIImage(response *TaskResponse, responseFormat string) *openai.ImageResponse {
+	imageResponse := openai.ImageResponse{
+		Created: helper.GetTimestamp(),
+	}
+
+	for _, data := range response.Output.Results {
+		var b64Json string
+		if responseFormat == "b64_json" {
+			// 读取 data.Url 的图片数据并转存到 b64Json
+			imageData, err := getImageData(data.Url)
+			if err != nil {
+				// 处理获取图片数据失败的情况
+				logger.SysError("getImageData Error getting image data: " + err.Error())
+				continue
+			}
+
+			// 将图片数据转为 Base64 编码的字符串
+			b64Json = Base64Encode(imageData)
+		} else {
+			// 如果 responseFormat 不是 "b64_json"，则直接使用 data.B64Image
+			b64Json = data.B64Image
+		}
+
+		imageResponse.Data = append(imageResponse.Data, openai.ImageData{
+			Url:           data.Url,
+			B64Json:       b64Json,
+			RevisedPrompt: "",
+		})
+	}
+	return &imageResponse
+}
+
+func getImageData(url string) ([]byte, error) {
+	response, err := http.Get(url)
+	if err != nil {
+		return nil, err
+	}
+	defer response.Body.Close()
+
+	imageData, err := io.ReadAll(response.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	return imageData, nil
+}
+
+func Base64Encode(data []byte) string {
+	b64Json := base64.StdEncoding.EncodeToString(data)
+	return b64Json
+}

relay/channel/ali/main.go

@@ -2,10 +2,7 @@ package ali
 
 import (
 	"bufio"
-	"encoding/base64"
 	"encoding/json"
-	"errors"
-	"fmt"
 	"github.com/gin-gonic/gin"
 	"github.com/songquanpeng/one-api/common"
 	"github.com/songquanpeng/one-api/common/helper"
@@ -15,7 +12,6 @@ import (
 	"io"
 	"net/http"
 	"strings"
-	"time"
 )
 
 // https://help.aliyun.com/document_detail/613695.html?spm=a2c4g.2399480.0.0.1adb778fAdzP9w#341800c0f8w0r
@@ -280,173 +276,3 @@ func Handler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *
 	_, err = c.Writer.Write(jsonResponse)
 	return nil, &fullTextResponse.Usage
 }
-
-func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) {
-	apiKey := c.Request.Header.Get("Authorization")
-	apiKey = strings.TrimPrefix(apiKey, "Bearer ")
-	responseFormat := c.GetString("response_format")
-
-	var aliTaskResponse TaskResponse
-	responseBody, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return openai.ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil
-	}
-	err = resp.Body.Close()
-	if err != nil {
-		return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
-	}
-	err = json.Unmarshal(responseBody, &aliTaskResponse)
-	if err != nil {
-		return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
-	}
-
-	if aliTaskResponse.Message != "" {
-		logger.SysError("aliAsyncTask err: " + string(responseBody))
-		return openai.ErrorWrapper(errors.New(aliTaskResponse.Message), "ali_async_task_failed", http.StatusInternalServerError), nil
-	}
-
-	aliResponse, err, _ := asyncTaskWait(aliTaskResponse.Output.TaskId, apiKey)
-	if err != nil {
-		return openai.ErrorWrapper(err, "ali_async_task_wait_failed", http.StatusInternalServerError), nil
-	}
-
-	if aliResponse.Output.TaskStatus != "SUCCEEDED" {
-		return &model.ErrorWithStatusCode{
-			Error: model.Error{
-				Message: aliResponse.Output.Message,
-				Type:    "ali_error",
-				Param:   "",
-				Code:    aliResponse.Output.Code,
-			},
-			StatusCode: resp.StatusCode,
-		}, nil
-	}
-
-	fullTextResponse := responseAli2OpenAIImage(aliResponse, responseFormat)
-	jsonResponse, err := json.Marshal(fullTextResponse)
-	if err != nil {
-		return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil
-	}
-	c.Writer.Header().Set("Content-Type", "application/json")
-	c.Writer.WriteHeader(resp.StatusCode)
-	_, err = c.Writer.Write(jsonResponse)
-	return nil, &fullTextResponse.Usage
-}
-
-func asyncTask(taskID string, key string) (*TaskResponse, error, []byte) {
-	url := fmt.Sprintf("https://dashscope.aliyuncs.com/api/v1/tasks/%s", taskID)
-
-	var aliResponse TaskResponse
-
-	req, err := http.NewRequest("GET", url, nil)
-	if err != nil {
-		return &aliResponse, err, nil
-	}
-
-	req.Header.Set("Authorization", "Bearer "+key)
-
-	client := &http.Client{}
-	resp, err := client.Do(req)
-	if err != nil {
-		logger.SysError("aliAsyncTask client.Do err: " + err.Error())
-		return &aliResponse, err, nil
-	}
-	defer resp.Body.Close()
-
-	responseBody, err := io.ReadAll(resp.Body)
-
-	var response TaskResponse
-	err = json.Unmarshal(responseBody, &response)
-	if err != nil {
-		logger.SysError("aliAsyncTask NewDecoder err: " + err.Error())
-		return &aliResponse, err, nil
-	}
-
-	return &response, nil, responseBody
-}
-
-func asyncTaskWait(taskID string, key string) (*TaskResponse, error, []byte) {
-	waitSeconds := 2
-	step := 0
-
-	var taskResponse TaskResponse
-	var responseBody []byte
-
-	for {
-		step++
-		rsp, err, body := asyncTask(taskID, key)
-		responseBody = body
-		if err != nil {
-			return &taskResponse, err, responseBody
-		}
-
-		if rsp.Output.TaskStatus == "" {
-			return &taskResponse, nil, responseBody
-		}
-
-		switch rsp.Output.TaskStatus {
-		case "FAILED":
-			fallthrough
-		case "CANCELED":
-			fallthrough
-		case "SUCCEEDED":
-			fallthrough
-		case "UNKNOWN":
-			return rsp, nil, responseBody
-		}
-
-		time.Sleep(time.Duration(waitSeconds) * time.Second)
-	}
-
-	return &taskResponse, nil, responseBody
-}
-
-func responseAli2OpenAIImage(response *TaskResponse, responseFormat string) *openai.ImageResponse {
-	imageRespones := openai.ImageResponse{}
-
-	for _, data := range response.Output.Results {
-		var b64Json string
-		if responseFormat == "b64_json" {
-			// 读取 data.Url 的图片数据并转存到 b64Json
-			imageData, err := getImageData(data.Url)
-			if err != nil {
-				// 处理获取图片数据失败的情况
-				logger.SysError("getImageData Error getting image data: " + err.Error())
-				continue
-			}
-
-			// 将图片数据转为 Base64 编码的字符串
-			b64Json = Base64Encode(imageData)
-		} else {
-			// 如果 responseFormat 不是 "b64_json"，则直接使用 data.B64Image
-			b64Json = data.B64Image
-		}
-
-		imageRespones.Data = append(imageRespones.Data, openai.ImageData{
-			Url:           data.Url,
-			B64Json:       b64Json,
-			RevisedPrompt: "",
-		})
-	}
-	return &imageRespones
-}
-
-func getImageData(url string) ([]byte, error) {
-	response, err := http.Get(url)
-	if err != nil {
-		return nil, err
-	}
-	defer response.Body.Close()
-
-	imageData, err := io.ReadAll(response.Body)
-	if err != nil {
-		return nil, err
-	}
-
-	return imageData, nil
-}
-
-func Base64Encode(data []byte) string {
-	b64Json := base64.StdEncoding.EncodeToString(data)
-	return b64Json
-}

relay/channel/openai/adaptor.go

@@ -92,7 +92,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.Rel
 	} else {
 		switch meta.Mode {
 		case constant.RelayModeImagesGenerations:
-			err, usage = ImageHandler(c, resp)
+			err, _ = ImageHandler(c, resp)
 		default:
 			err, usage = Handler(c, resp, meta.PromptTokens, meta.ActualModelName)
 		}

relay/channel/openai/main.go

@@ -181,5 +181,5 @@ func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCo
 	if err != nil {
 		return ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
 	}
-	return nil, &imageResponse.Usage
+	return nil, nil
 }

relay/channel/openai/model.go

@@ -117,9 +117,9 @@ type ImageData struct {
 }
 
 type ImageResponse struct {
-	Created     int         `json:"created"`
+	Created int64       `json:"created"`
 	Data    []ImageData `json:"data"`
-	model.Usage `json:"usage"`
+	//model.Usage `json:"usage"`
 }
 
 type ChatCompletionsStreamResponseChoice struct {