feat: 初步重构

relay/channel/openai/adaptor.go

@@ -0,0 +1,84 @@
+package openai
+
+import (
+	"errors"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"io"
+	"net/http"
+	"one-api/common"
+	"one-api/dto"
+	relaychannel "one-api/relay/channel"
+	relaycommon "one-api/relay/common"
+	"one-api/service"
+	"strings"
+)
+
+type Adaptor struct {
+}
+
+func (a *Adaptor) Init(info *relaycommon.RelayInfo, request dto.GeneralOpenAIRequest) {
+}
+
+func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
+	if info.ChannelType == common.ChannelTypeAzure {
+		// https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api
+		requestURL := strings.Split(info.RequestURLPath, "?")[0]
+		requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, info.ApiVersion)
+		task := strings.TrimPrefix(requestURL, "/v1/")
+		model_ := info.UpstreamModelName
+		model_ = strings.Replace(model_, ".", "", -1)
+		// https://github.com/songquanpeng/one-api/issues/67
+		model_ = strings.TrimSuffix(model_, "-0301")
+		model_ = strings.TrimSuffix(model_, "-0314")
+		model_ = strings.TrimSuffix(model_, "-0613")
+
+		requestURL = fmt.Sprintf("/openai/deployments/%s/%s", model_, task)
+		return relaycommon.GetFullRequestURL(info.BaseUrl, requestURL, info.ChannelType), nil
+	}
+	return relaycommon.GetFullRequestURL(info.BaseUrl, info.RequestURLPath, info.ChannelType), nil
+}
+
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+	relaychannel.SetupApiRequestHeader(info, c, req)
+	if info.ChannelType == common.ChannelTypeAzure {
+		req.Header.Set("api-key", info.ApiKey)
+		return nil
+	}
+	req.Header.Set("Authorization", "Bearer "+info.ApiKey)
+	if info.ChannelType == common.ChannelTypeOpenRouter {
+		req.Header.Set("HTTP-Referer", "https://github.com/songquanpeng/one-api")
+		req.Header.Set("X-Title", "One API")
+	}
+	return nil
+}
+
+func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *dto.GeneralOpenAIRequest) (any, error) {
+	if request == nil {
+		return nil, errors.New("request is nil")
+	}
+	return request, nil
+}
+
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+	return relaychannel.DoApiRequest(a, c, info, requestBody)
+}
+
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+	if info.IsStream {
+		var responseText string
+		err, responseText = openaiStreamHandler(c, resp, info.RelayMode)
+		usage = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
+	} else {
+		err, usage = openaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
+	}
+	return
+}
+
+func (a *Adaptor) GetModelList() []string {
+	return ModelList
+}
+
+func (a *Adaptor) GetChannelName() string {
+	return ChannelName
+}

relay/channel/openai/constant.go

@@ -0,0 +1,21 @@
+package openai
+
+var ModelList = []string{
+	"gpt-3.5-turbo", "gpt-3.5-turbo-0301", "gpt-3.5-turbo-0613", "gpt-3.5-turbo-1106", "gpt-3.5-turbo-0125",
+	"gpt-3.5-turbo-16k", "gpt-3.5-turbo-16k-0613",
+	"gpt-3.5-turbo-instruct",
+	"gpt-4", "gpt-4-0314", "gpt-4-0613", "gpt-4-1106-preview", "gpt-4-0125-preview",
+	"gpt-4-32k", "gpt-4-32k-0314", "gpt-4-32k-0613",
+	"gpt-4-turbo-preview",
+	"gpt-4-vision-preview",
+	"text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large",
+	"text-curie-001", "text-babbage-001", "text-ada-001", "text-davinci-002", "text-davinci-003",
+	"text-moderation-latest", "text-moderation-stable",
+	"text-davinci-edit-001",
+	"davinci-002", "babbage-002",
+	"dall-e-2", "dall-e-3",
+	"whisper-1",
+	"tts-1", "tts-1-1106", "tts-1-hd", "tts-1-hd-1106",
+}
+
+var ChannelName = "openai"

relay/channel/openai/relay-openai.go

@@ -0,0 +1,165 @@
+package openai
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/json"
+	"github.com/gin-gonic/gin"
+	"io"
+	"net/http"
+	"one-api/common"
+	"one-api/dto"
+	relayconstant "one-api/relay/constant"
+	"one-api/service"
+	"strings"
+	"sync"
+	"time"
+)
+
+func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*dto.OpenAIErrorWithStatusCode, string) {
+	var responseTextBuilder strings.Builder
+	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, 5)
+	stopChan := make(chan bool, 2)
+	defer close(stopChan)
+	defer close(dataChan)
+	var wg sync.WaitGroup
+
+	go func() {
+		wg.Add(1)
+		defer wg.Done()
+		var streamItems []string
+		for scanner.Scan() {
+			data := scanner.Text()
+			if len(data) < 6 { // ignore blank line or wrong format
+				continue
+			}
+			if data[:6] != "data: " && data[:6] != "[DONE]" {
+				continue
+			}
+			dataChan <- data
+			data = data[6:]
+			if !strings.HasPrefix(data, "[DONE]") {
+				streamItems = append(streamItems, data)
+			}
+		}
+		streamResp := "[" + strings.Join(streamItems, ",") + "]"
+		switch relayMode {
+		case relayconstant.RelayModeChatCompletions:
+			var streamResponses []dto.ChatCompletionsStreamResponseSimple
+			err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
+			if err != nil {
+				common.SysError("error unmarshalling stream response: " + err.Error())
+				return // just ignore the error
+			}
+			for _, streamResponse := range streamResponses {
+				for _, choice := range streamResponse.Choices {
+					responseTextBuilder.WriteString(choice.Delta.Content)
+				}
+			}
+		case relayconstant.RelayModeCompletions:
+			var streamResponses []dto.CompletionsStreamResponse
+			err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
+			if err != nil {
+				common.SysError("error unmarshalling stream response: " + err.Error())
+				return // just ignore the error
+			}
+			for _, streamResponse := range streamResponses {
+				for _, choice := range streamResponse.Choices {
+					responseTextBuilder.WriteString(choice.Text)
+				}
+			}
+		}
+		if len(dataChan) > 0 {
+			// wait data out
+			time.Sleep(2 * time.Second)
+		}
+		common.SafeSend(stopChan, true)
+	}()
+	service.SetEventStreamHeaders(c)
+	c.Stream(func(w io.Writer) bool {
+		select {
+		case data := <-dataChan:
+			if strings.HasPrefix(data, "data: [DONE]") {
+				data = data[:12]
+			}
+			// some implementations may add \r at the end of data
+			data = strings.TrimSuffix(data, "\r")
+			c.Render(-1, common.CustomEvent{Data: data})
+			return true
+		case <-stopChan:
+			return false
+		}
+	})
+	err := resp.Body.Close()
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
+	}
+	wg.Wait()
+	return nil, responseTextBuilder.String()
+}
+
+func openaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model string) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) {
+	var textResponse dto.TextResponse
+	responseBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = resp.Body.Close()
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = json.Unmarshal(responseBody, &textResponse)
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
+	}
+	if textResponse.Error.Type != "" {
+		return &dto.OpenAIErrorWithStatusCode{
+			OpenAIError: textResponse.Error,
+			StatusCode:  resp.StatusCode,
+		}, nil
+	}
+	// Reset response body
+	resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
+	// We shouldn't set the header before we parse the response body, because the parse part may fail.
+	// And then we will have to send an error response, but in this case, the header has already been set.
+	// So the httpClient will be confused by the response.
+	// For example, Postman will report error, and we cannot check the response at all.
+	for k, v := range resp.Header {
+		c.Writer.Header().Set(k, v[0])
+	}
+	c.Writer.WriteHeader(resp.StatusCode)
+	_, err = io.Copy(c.Writer, resp.Body)
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = resp.Body.Close()
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
+
+	if textResponse.Usage.TotalTokens == 0 {
+		completionTokens := 0
+		for _, choice := range textResponse.Choices {
+			completionTokens += service.CountTokenText(string(choice.Message.Content), model)
+		}
+		textResponse.Usage = dto.Usage{
+			PromptTokens:     promptTokens,
+			CompletionTokens: completionTokens,
+			TotalTokens:      promptTokens + completionTokens,
+		}
+	}
+	return nil, &textResponse.Usage
+}