♻️ refactor: split relay

common/client.go

@@ -0,0 +1,126 @@
+package common
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+)
+
+var HttpClient *http.Client
+
+func init() {
+	if RelayTimeout == 0 {
+		HttpClient = &http.Client{}
+	} else {
+		HttpClient = &http.Client{
+			Timeout: time.Duration(RelayTimeout) * time.Second,
+		}
+	}
+}
+
+type Client struct {
+	requestBuilder    RequestBuilder
+	createFormBuilder func(io.Writer) FormBuilder
+}
+
+func NewClient() *Client {
+	return &Client{
+		requestBuilder: NewRequestBuilder(),
+		createFormBuilder: func(body io.Writer) FormBuilder {
+			return NewFormBuilder(body)
+		},
+	}
+}
+
+type requestOptions struct {
+	body   any
+	header http.Header
+}
+
+type requestOption func(*requestOptions)
+
+func WithBody(body any) requestOption {
+	return func(args *requestOptions) {
+		args.body = body
+	}
+}
+
+func WithHeader(header map[string]string) requestOption {
+	return func(args *requestOptions) {
+		for k, v := range header {
+			args.header.Set(k, v)
+		}
+	}
+}
+
+type RequestError struct {
+	HTTPStatusCode int
+	Err            error
+}
+
+func (c *Client) NewRequest(method, url string, setters ...requestOption) (*http.Request, error) {
+	// Default Options
+	args := &requestOptions{
+		body:   nil,
+		header: make(http.Header),
+	}
+	for _, setter := range setters {
+		setter(args)
+	}
+	req, err := c.requestBuilder.Build(method, url, args.body, args.header)
+	if err != nil {
+		return nil, err
+	}
+
+	return req, nil
+}
+
+func (c *Client) SendRequest(req *http.Request, response any) error {
+
+	// 发送请求
+	resp, err := HttpClient.Do(req)
+	if err != nil {
+		return err
+	}
+
+	defer resp.Body.Close()
+
+	// 处理响应
+	if IsFailureStatusCode(resp) {
+		return fmt.Errorf("status code: %d", resp.StatusCode)
+	}
+
+	// 解析响应
+	err = DecodeResponse(resp.Body, response)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func IsFailureStatusCode(resp *http.Response) bool {
+	return resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusBadRequest
+}
+
+func DecodeResponse(body io.Reader, v any) error {
+	if v == nil {
+		return nil
+	}
+
+	if result, ok := v.(*string); ok {
+		return DecodeString(body, result)
+	}
+	return json.NewDecoder(body).Decode(v)
+}
+
+func DecodeString(body io.Reader, output *string) error {
+	b, err := io.ReadAll(body)
+	if err != nil {
+		return err
+	}
+	*output = string(b)
+	return nil
+}

common/form_builder.go

@@ -0,0 +1,65 @@
+package common
+
+import (
+	"fmt"
+	"io"
+	"mime/multipart"
+	"os"
+	"path"
+)
+
+type FormBuilder interface {
+	CreateFormFile(fieldname string, file *os.File) error
+	CreateFormFileReader(fieldname string, r io.Reader, filename string) error
+	WriteField(fieldname, value string) error
+	Close() error
+	FormDataContentType() string
+}
+
+type DefaultFormBuilder struct {
+	writer *multipart.Writer
+}
+
+func NewFormBuilder(body io.Writer) *DefaultFormBuilder {
+	return &DefaultFormBuilder{
+		writer: multipart.NewWriter(body),
+	}
+}
+
+func (fb *DefaultFormBuilder) CreateFormFile(fieldname string, file *os.File) error {
+	return fb.createFormFile(fieldname, file, file.Name())
+}
+
+func (fb *DefaultFormBuilder) CreateFormFileReader(fieldname string, r io.Reader, filename string) error {
+	return fb.createFormFile(fieldname, r, path.Base(filename))
+}
+
+func (fb *DefaultFormBuilder) createFormFile(fieldname string, r io.Reader, filename string) error {
+	if filename == "" {
+		return fmt.Errorf("filename cannot be empty")
+	}
+
+	fieldWriter, err := fb.writer.CreateFormFile(fieldname, filename)
+	if err != nil {
+		return err
+	}
+
+	_, err = io.Copy(fieldWriter, r)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (fb *DefaultFormBuilder) WriteField(fieldname, value string) error {
+	return fb.writer.WriteField(fieldname, value)
+}
+
+func (fb *DefaultFormBuilder) Close() error {
+	return fb.writer.Close()
+}
+
+func (fb *DefaultFormBuilder) FormDataContentType() string {
+	return fb.writer.FormDataContentType()
+}

common/marshaller.go

@@ -0,0 +1,15 @@
+package common
+
+import (
+	"encoding/json"
+)
+
+type Marshaller interface {
+	Marshal(value any) ([]byte, error)
+}
+
+type JSONMarshaller struct{}
+
+func (jm *JSONMarshaller) Marshal(value any) ([]byte, error) {
+	return json.Marshal(value)
+}

common/quota.go

@@ -0,0 +1,59 @@
+package common
+
+// type Quota struct {
+// 	ModelName  string
+// 	ModelRatio float64
+// 	GroupRatio float64
+// 	Ratio      float64
+// 	UserQuota  int
+// }
+
+// func CreateQuota(modelName string, userQuota int, group string) *Quota {
+// 	modelRatio := GetModelRatio(modelName)
+// 	groupRatio := GetGroupRatio(group)
+
+// 	return &Quota{
+// 		ModelName:  modelName,
+// 		ModelRatio: modelRatio,
+// 		GroupRatio: groupRatio,
+// 		Ratio:      modelRatio * groupRatio,
+// 		UserQuota:  userQuota,
+// 	}
+// }
+
+// func (q *Quota) getTokenNum(tokenEncoder *tiktoken.Tiktoken, text string) int {
+// 	if ApproximateTokenEnabled {
+// 		return int(float64(len(text)) * 0.38)
+// 	}
+// 	return len(tokenEncoder.Encode(text, nil, nil))
+// }
+
+// func (q *Quota) CountTokenMessages(messages []Message, model string) int {
+// 	tokenEncoder := q.getTokenEncoder(model)
+// 	// Reference:
+// 	// https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
+// 	// https://github.com/pkoukk/tiktoken-go/issues/6
+// 	//
+// 	// Every message follows <|start|>{role/name}\n{content}<|end|>\n
+// 	var tokensPerMessage int
+// 	var tokensPerName int
+// 	if model == "gpt-3.5-turbo-0301" {
+// 		tokensPerMessage = 4
+// 		tokensPerName = -1 // If there's a name, the role is omitted
+// 	} else {
+// 		tokensPerMessage = 3
+// 		tokensPerName = 1
+// 	}
+// 	tokenNum := 0
+// 	for _, message := range messages {
+// 		tokenNum += tokensPerMessage
+// 		tokenNum += q.getTokenNum(tokenEncoder, message.StringContent())
+// 		tokenNum += q.getTokenNum(tokenEncoder, message.Role)
+// 		if message.Name != nil {
+// 			tokenNum += tokensPerName
+// 			tokenNum += q.getTokenNum(tokenEncoder, *message.Name)
+// 		}
+// 	}
+// 	tokenNum += 3 // Every reply is primed with <|start|>assistant<|message|>
+// 	return tokenNum
+// }

common/request_builder.go

@@ -0,0 +1,50 @@
+package common
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+)
+
+type RequestBuilder interface {
+	Build(method, url string, body any, header http.Header) (*http.Request, error)
+}
+
+type HTTPRequestBuilder struct {
+	marshaller Marshaller
+}
+
+func NewRequestBuilder() *HTTPRequestBuilder {
+	return &HTTPRequestBuilder{
+		marshaller: &JSONMarshaller{},
+	}
+}
+
+func (b *HTTPRequestBuilder) Build(
+	method string,
+	url string,
+	body any,
+	header http.Header,
+) (req *http.Request, err error) {
+	var bodyReader io.Reader
+	if body != nil {
+		if v, ok := body.(io.Reader); ok {
+			bodyReader = v
+		} else {
+			var reqBytes []byte
+			reqBytes, err = b.marshaller.Marshal(body)
+			if err != nil {
+				return
+			}
+			bodyReader = bytes.NewBuffer(reqBytes)
+		}
+	}
+	req, err = http.NewRequest(method, url, bodyReader)
+	if err != nil {
+		return
+	}
+	if header != nil {
+		req.Header = header
+	}
+	return
+}

common/token.go

@@ -0,0 +1,109 @@
+package common
+
+import (
+	"fmt"
+	"strings"
+
+	"one-api/types"
+
+	"github.com/pkoukk/tiktoken-go"
+)
+
+var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
+var defaultTokenEncoder *tiktoken.Tiktoken
+
+func InitTokenEncoders() {
+	SysLog("initializing token encoders")
+	gpt35TokenEncoder, err := tiktoken.EncodingForModel("gpt-3.5-turbo")
+	if err != nil {
+		FatalLog(fmt.Sprintf("failed to get gpt-3.5-turbo token encoder: %s", err.Error()))
+	}
+	defaultTokenEncoder = gpt35TokenEncoder
+	gpt4TokenEncoder, err := tiktoken.EncodingForModel("gpt-4")
+	if err != nil {
+		FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
+	}
+	for model, _ := range ModelRatio {
+		if strings.HasPrefix(model, "gpt-3.5") {
+			tokenEncoderMap[model] = gpt35TokenEncoder
+		} else if strings.HasPrefix(model, "gpt-4") {
+			tokenEncoderMap[model] = gpt4TokenEncoder
+		} else {
+			tokenEncoderMap[model] = nil
+		}
+	}
+	SysLog("token encoders initialized")
+}
+
+func getTokenEncoder(model string) *tiktoken.Tiktoken {
+	tokenEncoder, ok := tokenEncoderMap[model]
+	if ok && tokenEncoder != nil {
+		return tokenEncoder
+	}
+	if ok {
+		tokenEncoder, err := tiktoken.EncodingForModel(model)
+		if err != nil {
+			SysError(fmt.Sprintf("failed to get token encoder for model %s: %s, using encoder for gpt-3.5-turbo", model, err.Error()))
+			tokenEncoder = defaultTokenEncoder
+		}
+		tokenEncoderMap[model] = tokenEncoder
+		return tokenEncoder
+	}
+	return defaultTokenEncoder
+}
+
+func getTokenNum(tokenEncoder *tiktoken.Tiktoken, text string) int {
+	if ApproximateTokenEnabled {
+		return int(float64(len(text)) * 0.38)
+	}
+	return len(tokenEncoder.Encode(text, nil, nil))
+}
+
+func CountTokenMessages(messages []types.ChatCompletionMessage, model string) int {
+	tokenEncoder := getTokenEncoder(model)
+	// Reference:
+	// https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
+	// https://github.com/pkoukk/tiktoken-go/issues/6
+	//
+	// Every message follows <|start|>{role/name}\n{content}<|end|>\n
+	var tokensPerMessage int
+	var tokensPerName int
+	if model == "gpt-3.5-turbo-0301" {
+		tokensPerMessage = 4
+		tokensPerName = -1 // If there's a name, the role is omitted
+	} else {
+		tokensPerMessage = 3
+		tokensPerName = 1
+	}
+	tokenNum := 0
+	for _, message := range messages {
+		tokenNum += tokensPerMessage
+		tokenNum += getTokenNum(tokenEncoder, message.StringContent())
+		tokenNum += getTokenNum(tokenEncoder, message.Role)
+		if message.Name != nil {
+			tokenNum += tokensPerName
+			tokenNum += getTokenNum(tokenEncoder, *message.Name)
+		}
+	}
+	tokenNum += 3 // Every reply is primed with <|start|>assistant<|message|>
+	return tokenNum
+}
+
+func CountTokenInput(input any, model string) int {
+	switch input.(type) {
+	case string:
+		return CountTokenText(input.(string), model)
+	case []string:
+		text := ""
+		for _, s := range input.([]string) {
+			text += s
+		}
+		return CountTokenText(text, model)
+	}
+	return 0
+}
+
+func CountTokenText(text string, model string) int {
+	tokenEncoder := getTokenEncoder(model)
+	return getTokenNum(tokenEncoder, text)
+}