feat: support baidu's models now (close #286)

controller/relay-baidu.go

@@ -0,0 +1,203 @@
+package controller
+
+import (
+	"bufio"
+	"encoding/json"
+	"github.com/gin-gonic/gin"
+	"io"
+	"net/http"
+	"one-api/common"
+	"strings"
+)
+
+// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/flfmc9do2
+
+type BaiduTokenResponse struct {
+	RefreshToken  string `json:"refresh_token"`
+	ExpiresIn     int    `json:"expires_in"`
+	SessionKey    string `json:"session_key"`
+	AccessToken   string `json:"access_token"`
+	Scope         string `json:"scope"`
+	SessionSecret string `json:"session_secret"`
+}
+
+type BaiduMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type BaiduChatRequest struct {
+	Messages []BaiduMessage `json:"messages"`
+	Stream   bool           `json:"stream"`
+	UserId   string         `json:"user_id,omitempty"`
+}
+
+type BaiduError struct {
+	ErrorCode int    `json:"error_code"`
+	ErrorMsg  string `json:"error_msg"`
+}
+
+type BaiduChatResponse struct {
+	Id               string `json:"id"`
+	Object           string `json:"object"`
+	Created          int64  `json:"created"`
+	Result           string `json:"result"`
+	IsTruncated      bool   `json:"is_truncated"`
+	NeedClearHistory bool   `json:"need_clear_history"`
+	Usage            Usage  `json:"usage"`
+	BaiduError
+}
+
+type BaiduChatStreamResponse struct {
+	BaiduChatResponse
+	SentenceId int  `json:"sentence_id"`
+	IsEnd      bool `json:"is_end"`
+}
+
+func requestOpenAI2Baidu(request GeneralOpenAIRequest) *BaiduChatRequest {
+	messages := make([]BaiduMessage, 0, len(request.Messages))
+	for _, message := range request.Messages {
+		messages = append(messages, BaiduMessage{
+			Role:    message.Role,
+			Content: message.Content,
+		})
+	}
+	return &BaiduChatRequest{
+		Messages: messages,
+		Stream:   request.Stream,
+	}
+}
+
+func responseBaidu2OpenAI(response *BaiduChatResponse) *OpenAITextResponse {
+	choice := OpenAITextResponseChoice{
+		Index: 0,
+		Message: Message{
+			Role:    "assistant",
+			Content: response.Result,
+		},
+		FinishReason: "stop",
+	}
+	fullTextResponse := OpenAITextResponse{
+		Id:      response.Id,
+		Object:  "chat.completion",
+		Created: response.Created,
+		Choices: []OpenAITextResponseChoice{choice},
+		Usage:   response.Usage,
+	}
+	return &fullTextResponse
+}
+
+func streamResponseBaidu2OpenAI(baiduResponse *BaiduChatStreamResponse) *ChatCompletionsStreamResponse {
+	var choice ChatCompletionsStreamResponseChoice
+	choice.Delta.Content = baiduResponse.Result
+	choice.FinishReason = "stop"
+	response := ChatCompletionsStreamResponse{
+		Id:      baiduResponse.Id,
+		Object:  "chat.completion.chunk",
+		Created: baiduResponse.Created,
+		Model:   "ernie-bot",
+		Choices: []ChatCompletionsStreamResponseChoice{choice},
+	}
+	return &response
+}
+
+func baiduStreamHandler(c *gin.Context, resp *http.Response) (*OpenAIErrorWithStatusCode, *Usage) {
+	var usage Usage
+	scanner := bufio.NewScanner(resp.Body)
+	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
+		if atEOF && len(data) == 0 {
+			return 0, nil, nil
+		}
+		if i := strings.Index(string(data), "\n"); i >= 0 {
+			return i + 1, data[0:i], nil
+		}
+		if atEOF {
+			return len(data), data, nil
+		}
+		return 0, nil, nil
+	})
+	dataChan := make(chan string)
+	stopChan := make(chan bool)
+	go func() {
+		for scanner.Scan() {
+			data := scanner.Text()
+			if len(data) < 6 { // ignore blank line or wrong format
+				continue
+			}
+			data = data[6:]
+			dataChan <- data
+		}
+		stopChan <- true
+	}()
+	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	c.Writer.Header().Set("Cache-Control", "no-cache")
+	c.Writer.Header().Set("Connection", "keep-alive")
+	c.Writer.Header().Set("Transfer-Encoding", "chunked")
+	c.Writer.Header().Set("X-Accel-Buffering", "no")
+	c.Stream(func(w io.Writer) bool {
+		select {
+		case data := <-dataChan:
+			var baiduResponse BaiduChatStreamResponse
+			err := json.Unmarshal([]byte(data), &baiduResponse)
+			if err != nil {
+				common.SysError("error unmarshalling stream response: " + err.Error())
+				return true
+			}
+			usage.PromptTokens += baiduResponse.Usage.PromptTokens
+			usage.CompletionTokens += baiduResponse.Usage.CompletionTokens
+			usage.TotalTokens += baiduResponse.Usage.TotalTokens
+			response := streamResponseBaidu2OpenAI(&baiduResponse)
+			jsonResponse, err := json.Marshal(response)
+			if err != nil {
+				common.SysError("error marshalling stream response: " + err.Error())
+				return true
+			}
+			c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)})
+			return true
+		case <-stopChan:
+			c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
+			return false
+		}
+	})
+	err := resp.Body.Close()
+	if err != nil {
+		return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
+	return nil, &usage
+}
+
+func baiduHandler(c *gin.Context, resp *http.Response) (*OpenAIErrorWithStatusCode, *Usage) {
+	var baiduResponse BaiduChatResponse
+	responseBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return errorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = resp.Body.Close()
+	if err != nil {
+		return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = json.Unmarshal(responseBody, &baiduResponse)
+	if err != nil {
+		return errorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
+	}
+	if baiduResponse.ErrorMsg != "" {
+		return &OpenAIErrorWithStatusCode{
+			OpenAIError: OpenAIError{
+				Message: baiduResponse.ErrorMsg,
+				Type:    "baidu_error",
+				Param:   "",
+				Code:    baiduResponse.ErrorCode,
+			},
+			StatusCode: resp.StatusCode,
+		}, nil
+	}
+	fullTextResponse := responseBaidu2OpenAI(&baiduResponse)
+	jsonResponse, err := json.Marshal(fullTextResponse)
+	if err != nil {
+		return 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
+}