Merge branch 'main' into pr/Laisky/25

relay/controller/text.go

@@ -0,0 +1,144 @@
+package controller
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/ctxkey"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/model"
+	"github.com/songquanpeng/one-api/relay"
+	"github.com/songquanpeng/one-api/relay/adaptor"
+	"github.com/songquanpeng/one-api/relay/adaptor/openai"
+	"github.com/songquanpeng/one-api/relay/apitype"
+	"github.com/songquanpeng/one-api/relay/billing"
+	billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
+	"github.com/songquanpeng/one-api/relay/channeltype"
+	"github.com/songquanpeng/one-api/relay/meta"
+	relaymodel "github.com/songquanpeng/one-api/relay/model"
+)
+
+func RelayTextHelper(c *gin.Context) *relaymodel.ErrorWithStatusCode {
+	ctx := c.Request.Context()
+	meta := meta.GetByContext(c)
+	// get & validate textRequest
+	textRequest, err := getAndValidateTextRequest(c, meta.Mode)
+	if err != nil {
+		logger.Errorf(ctx, "getAndValidateTextRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "invalid_text_request", http.StatusBadRequest)
+	}
+	meta.IsStream = textRequest.Stream
+
+	// map model name
+	meta.OriginModelName = textRequest.Model
+	textRequest.Model, _ = getMappedModelName(textRequest.Model, meta.ModelMapping)
+	meta.ActualModelName = textRequest.Model
+	// set system prompt if not empty
+	systemPromptReset := setSystemPrompt(ctx, textRequest, meta.SystemPrompt)
+	// get model ratio & group ratio
+	modelRatio := billingratio.GetModelRatio(textRequest.Model, meta.ChannelType)
+	// groupRatio := billingratio.GetGroupRatio(meta.Group)
+	groupRatio := c.GetFloat64(ctxkey.ChannelRatio)
+
+	ratio := modelRatio * groupRatio
+	// pre-consume quota
+	promptTokens := getPromptTokens(c.Request.Context(), textRequest, meta.Mode)
+	meta.PromptTokens = promptTokens
+	preConsumedQuota, bizErr := preConsumeQuota(c, textRequest, promptTokens, ratio, meta)
+	if bizErr != nil {
+		logger.Warnf(ctx, "preConsumeQuota failed: %+v", *bizErr)
+		return bizErr
+	}
+
+	adaptor := relay.GetAdaptor(meta.APIType)
+	if adaptor == nil {
+		return openai.ErrorWrapper(errors.Errorf("invalid api type: %d", meta.APIType), "invalid_api_type", http.StatusBadRequest)
+	}
+	adaptor.Init(meta)
+
+	// get request body
+	requestBody, err := getRequestBody(c, meta, textRequest, adaptor)
+	if err != nil {
+		return openai.ErrorWrapper(err, "convert_request_failed", http.StatusInternalServerError)
+	}
+
+	// for debug
+	requestBodyBytes, _ := io.ReadAll(requestBody)
+	requestBody = bytes.NewBuffer(requestBodyBytes)
+
+	// do request
+	resp, err := adaptor.DoRequest(c, meta, requestBody)
+	if err != nil {
+		logger.Errorf(ctx, "DoRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
+	}
+	if isErrorHappened(meta, resp) {
+		billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
+		return RelayErrorHandler(resp)
+	}
+
+	// do response
+	usage, respErr := adaptor.DoResponse(c, resp, meta)
+	if respErr != nil {
+		logger.Errorf(ctx, "respErr is not nil: %+v", respErr)
+		billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
+		return respErr
+	}
+
+	// post-consume quota
+	quotaId := c.GetInt(ctxkey.Id)
+	requestId := c.GetString(ctxkey.RequestId)
+	go func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cancel()
+
+		quota := postConsumeQuota(ctx, usage, meta, textRequest, ratio, preConsumedQuota, modelRatio, groupRatio, systemPromptReset)
+
+		// also update user request cost
+		if quota != 0 {
+			docu := model.NewUserRequestCost(
+				quotaId,
+				requestId,
+				quota,
+			)
+			if err = docu.Insert(); err != nil {
+				logger.Errorf(ctx, "insert user request cost failed: %+v", err)
+			}
+		}
+	}()
+
+	return nil
+}
+
+func getRequestBody(c *gin.Context, meta *meta.Meta, textRequest *relaymodel.GeneralOpenAIRequest, adaptor adaptor.Adaptor) (io.Reader, error) {
+	if !config.EnforceIncludeUsage &&
+		meta.APIType == apitype.OpenAI &&
+		meta.OriginModelName == meta.ActualModelName &&
+		meta.ChannelType != channeltype.OpenAI && // openai also need to convert request
+		meta.ChannelType != channeltype.Baichuan {
+		return c.Request.Body, nil
+	}
+
+	// get request body
+	var requestBody io.Reader
+	convertedRequest, err := adaptor.ConvertRequest(c, meta.Mode, textRequest)
+	if err != nil {
+		logger.Debugf(c.Request.Context(), "converted request failed: %s\n", err.Error())
+		return nil, err
+	}
+	jsonData, err := json.Marshal(convertedRequest)
+	if err != nil {
+		logger.Debugf(c.Request.Context(), "converted request json_marshal_failed: %s\n", err.Error())
+		return nil, err
+	}
+	logger.Debugf(c.Request.Context(), "converted request: \n%s", string(jsonData))
+	requestBody = bytes.NewBuffer(jsonData)
+	return requestBody, nil
+}