Skip to content

phsl / new-api

Go to a project
Commit eece0732 by Calcium-Ion Committed by GitHub
Options

Merge pull request #1642 from QuantumNous/fix-retry-and-rerank 

fix: retry requeset body incorrect and fix rerank
parents 54e03c42 13301d85
Showing with 243 additions and 116 deletions
common/copy.go 0 → 100644
package common
import (
"fmt"
"github.com/antlabs/pcopy"
)
func DeepCopy[T any](src *T) (*T, error) {
if src == nil {
return nil, fmt.Errorf("copy source cannot be nil")
}
var dst T
err := pcopy.Copy(&dst, src)
if err != nil {
return nil, err
}
if &dst == nil {
return nil, fmt.Errorf("copy result cannot be nil")
}
return &dst, nil
}
common/gin.go
......@@ -2,12 +2,13 @@ package common
import (
"bytes"
"github.com/gin-gonic/gin"
"io"
"net/http"
"one-api/constant"
"strings"
"time"
"github.com/gin-gonic/gin"
)
const KeyRequestBody = "key_request_body"
......
dto/audio.go
......@@ -26,6 +26,12 @@ func (r *AudioRequest) IsStream(c *gin.Context) bool {
return false
}
func (r *AudioRequest) SetModelName(modelName string) {
if modelName != "" {
r.Model = modelName
}
}
type AudioResponse struct {
Text string `json:"text"`
}
......
dto/claude.go
......@@ -321,8 +321,14 @@ func (c *ClaudeRequest) GetTokenCountMeta() *types.TokenCountMeta {
return &tokenCountMeta
}
func (claudeRequest *ClaudeRequest) IsStream(c *gin.Context) bool {
return claudeRequest.Stream
func (c *ClaudeRequest) IsStream(ctx *gin.Context) bool {
return c.Stream
}
func (c *ClaudeRequest) SetModelName(modelName string) {
if modelName != "" {
c.Model = modelName
}
}
func (c *ClaudeRequest) SearchToolNameByToolCallId(toolCallId string) string {
......
dto/embedding.go
......@@ -48,6 +48,12 @@ func (r *EmbeddingRequest) IsStream(c *gin.Context) bool {
return false
}
func (r *EmbeddingRequest) SetModelName(modelName string) {
if modelName != "" {
r.Model = modelName
}
}
func (r *EmbeddingRequest) ParseInput() []string {
if r.Input == nil {
return make([]string, 0)
......
dto/gemini.go
......@@ -73,6 +73,10 @@ func (r *GeminiChatRequest) IsStream(c *gin.Context) bool {
return false
}
func (r *GeminiChatRequest) SetModelName(modelName string) {
// GeminiChatRequest does not have a model field, so this method does nothing.
}
func (r *GeminiChatRequest) GetTools() []GeminiChatTool {
var tools []GeminiChatTool
if strings.HasSuffix(string(r.Tools), "[") {
......@@ -312,10 +316,61 @@ type GeminiEmbeddingRequest struct {
OutputDimensionality int `json:"outputDimensionality,omitempty"`
}
func (r *GeminiEmbeddingRequest) IsStream(c *gin.Context) bool {
// Gemini embedding requests are not streamed
return false
}
func (r *GeminiEmbeddingRequest) GetTokenCountMeta() *types.TokenCountMeta {
var inputTexts []string
for _, part := range r.Content.Parts {
if part.Text != "" {
inputTexts = append(inputTexts, part.Text)
}
}
inputText := strings.Join(inputTexts, "\n")
return &types.TokenCountMeta{
CombineText: inputText,
}
}
func (r *GeminiEmbeddingRequest) SetModelName(modelName string) {
if modelName != "" {
r.Model = modelName
}
}
type GeminiBatchEmbeddingRequest struct {
Requests []*GeminiEmbeddingRequest `json:"requests"`
}
func (r *GeminiBatchEmbeddingRequest) IsStream(c *gin.Context) bool {
// Gemini batch embedding requests are not streamed
return false
}
func (r *GeminiBatchEmbeddingRequest) GetTokenCountMeta() *types.TokenCountMeta {
var inputTexts []string
for _, request := range r.Requests {
meta := request.GetTokenCountMeta()
if meta != nil && meta.CombineText != "" {
inputTexts = append(inputTexts, meta.CombineText)
}
}
inputText := strings.Join(inputTexts, "\n")
return &types.TokenCountMeta{
CombineText: inputText,
}
}
func (r *GeminiBatchEmbeddingRequest) SetModelName(modelName string) {
if modelName != "" {
for _, req := range r.Requests {
req.SetModelName(modelName)
}
}
}
type GeminiEmbeddingResponse struct {
Embedding ContentEmbedding `json:"embedding"`
}
......
dto/openai_image.go
......@@ -12,10 +12,10 @@ type ImageRequest struct {
Model string `json:"model"`
Prompt string `json:"prompt" binding:"required"`
N uint `json:"n,omitempty"`
Size string `json:"size,omitempty"`
Quality string `json:"quality,omitempty"`
ResponseFormat string `json:"response_format,omitempty"`
Style json.RawMessage `json:"style,omitempty"`
Size string `json:"size,omitempty"`
Quality string `json:"quality,omitempty"`
ResponseFormat string `json:"response_format,omitempty"`
Style json.RawMessage `json:"style,omitempty"`
User json.RawMessage `json:"user,omitempty"`
ExtraFields json.RawMessage `json:"extra_fields,omitempty"`
Background json.RawMessage `json:"background,omitempty"`
......@@ -63,6 +63,12 @@ func (i *ImageRequest) IsStream(c *gin.Context) bool {
return false
}
func (i *ImageRequest) SetModelName(modelName string) {
if modelName != "" {
i.Model = modelName
}
}
type ImageResponse struct {
Data []ImageData `json:"data"`
Created int64 `json:"created"`
......
dto/openai_request.go
......@@ -183,6 +183,12 @@ func (r *GeneralOpenAIRequest) IsStream(c *gin.Context) bool {
return r.Stream
}
func (r *GeneralOpenAIRequest) SetModelName(modelName string) {
if modelName != "" {
r.Model = modelName
}
}
func (r *GeneralOpenAIRequest) ToMap() map[string]any {
result := make(map[string]any)
data, _ := common.Marshal(r)
......@@ -841,6 +847,12 @@ func (r *OpenAIResponsesRequest) IsStream(c *gin.Context) bool {
return r.Stream
}
func (r *OpenAIResponsesRequest) SetModelName(modelName string) {
if modelName != "" {
r.Model = modelName
}
}
type Reasoning struct {
Effort string `json:"effort,omitempty"`
Summary string `json:"summary,omitempty"`
......
dto/request_common.go
......@@ -8,6 +8,7 @@ import (
type Request interface {
GetTokenCountMeta() *types.TokenCountMeta
IsStream(c *gin.Context) bool
SetModelName(modelName string)
}
type BaseRequest struct {
......@@ -18,7 +19,7 @@ func (b *BaseRequest) GetTokenCountMeta() *types.TokenCountMeta {
TokenType: types.TokenTypeTokenizer,
}
}
func (b *BaseRequest) IsStream(c *gin.Context) bool {
return false
}
func (b *BaseRequest) SetModelName(modelName string) {}
dto/rerank.go
......@@ -37,6 +37,12 @@ func (r *RerankRequest) GetTokenCountMeta() *types.TokenCountMeta {
}
}
func (r *RerankRequest) SetModelName(modelName string) {
if modelName != "" {
r.Model = modelName
}
}
func (r *RerankRequest) GetReturnDocuments() bool {
if r.ReturnDocuments == nil {
return false
......
go.mod
......@@ -44,7 +44,11 @@ require (
)
require (
github.com/Masterminds/goutils v1.1.1 // indirect
github.com/Masterminds/semver/v3 v3.2.0 // indirect
github.com/Masterminds/sprig/v3 v3.2.3 // indirect
github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 // indirect
github.com/antlabs/pcopy v0.1.5 // indirect
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.0 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.2 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.2 // indirect
......@@ -69,6 +73,8 @@ require (
github.com/gorilla/context v1.1.1 // indirect
github.com/gorilla/securecookie v1.1.1 // indirect
github.com/gorilla/sessions v1.2.1 // indirect
github.com/huandu/xstrings v1.3.3 // indirect
github.com/imdario/mergo v0.3.11 // indirect
github.com/jackc/pgpassfile v1.0.0 // indirect
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
github.com/jackc/pgx/v5 v5.7.1 // indirect
......@@ -79,11 +85,14 @@ require (
github.com/klauspost/cpuid/v2 v2.2.9 // indirect
github.com/leodido/go-urn v1.4.0 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/mitchellh/copystructure v1.0.0 // indirect
github.com/mitchellh/mapstructure v1.5.0 // indirect
github.com/mitchellh/reflectwalk v1.0.0 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/pelletier/go-toml/v2 v2.2.1 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/spf13/cast v1.3.1 // indirect
github.com/tidwall/match v1.1.1 // indirect
github.com/tidwall/pretty v1.2.0 // indirect
github.com/tklauser/go-sysconf v0.3.12 // indirect
......
relay/audio_handler.go
......@@ -4,6 +4,7 @@ import (
"errors"
"fmt"
"net/http"
"one-api/common"
"one-api/dto"
relaycommon "one-api/relay/common"
"one-api/relay/helper"
......@@ -16,12 +17,17 @@ import (
func AudioHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
audioRequest, ok := info.Request.(*dto.AudioRequest)
audioReq, ok := info.Request.(*dto.AudioRequest)
if !ok {
return types.NewError(errors.New("invalid request type"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
err := helper.ModelMappedHelper(c, info, audioRequest)
request, err := common.DeepCopy(audioReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to AudioRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
......@@ -32,7 +38,7 @@ func AudioHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
}
adaptor.Init(info)
ioReader, err := adaptor.ConvertAudioRequest(c, info, *audioRequest)
ioReader, err := adaptor.ConvertAudioRequest(c, info, *request)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
......
relay/claude_handler.go
......@@ -21,13 +21,18 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
info.InitChannelMeta(c)
textRequest, ok := info.Request.(*dto.ClaudeRequest)
claudeReq, ok := info.Request.(*dto.ClaudeRequest)
if !ok {
common.FatalLog(fmt.Sprintf("invalid request type, expected *dto.ClaudeRequest, got %T", info.Request))
return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected *dto.ClaudeRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
err := helper.ModelMappedHelper(c, info, textRequest)
request, err := common.DeepCopy(claudeReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to ClaudeRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
......@@ -38,30 +43,30 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
adaptor.Init(info)
if textRequest.MaxTokens == 0 {
textRequest.MaxTokens = uint(model_setting.GetClaudeSettings().GetDefaultMaxTokens(textRequest.Model))
if request.MaxTokens == 0 {
request.MaxTokens = uint(model_setting.GetClaudeSettings().GetDefaultMaxTokens(request.Model))
}
if model_setting.GetClaudeSettings().ThinkingAdapterEnabled &&
strings.HasSuffix(textRequest.Model, "-thinking") {
if textRequest.Thinking == nil {
strings.HasSuffix(request.Model, "-thinking") {
if request.Thinking == nil {
// 因为BudgetTokens 必须大于1024
if textRequest.MaxTokens < 1280 {
textRequest.MaxTokens = 1280
if request.MaxTokens < 1280 {
request.MaxTokens = 1280
}
// BudgetTokens 为 max_tokens 的 80%
textRequest.Thinking = &dto.Thinking{
request.Thinking = &dto.Thinking{
Type: "enabled",
BudgetTokens: common.GetPointer[int](int(float64(textRequest.MaxTokens) * model_setting.GetClaudeSettings().ThinkingAdapterBudgetTokensPercentage)),
BudgetTokens: common.GetPointer[int](int(float64(request.MaxTokens) * model_setting.GetClaudeSettings().ThinkingAdapterBudgetTokensPercentage)),
}
// TODO: 临时处理
// https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#important-considerations-when-using-extended-thinking
textRequest.TopP = 0
textRequest.Temperature = common.GetPointer[float64](1.0)
request.TopP = 0
request.Temperature = common.GetPointer[float64](1.0)
}
textRequest.Model = strings.TrimSuffix(textRequest.Model, "-thinking")
info.UpstreamModelName = textRequest.Model
request.Model = strings.TrimSuffix(request.Model, "-thinking")
info.UpstreamModelName = request.Model
}
var requestBody io.Reader
......@@ -72,7 +77,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
requestBody = bytes.NewBuffer(body)
} else {
convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, textRequest)
convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
......
relay/common/relay_info.go
......@@ -158,7 +158,14 @@ func (info *RelayInfo) InitChannelMeta(c *gin.Context) {
if streamSupportedChannels[channelMeta.ChannelType] {
channelMeta.SupportStreamOptions = true
}
info.ChannelMeta = channelMeta
// reset some fields based on channel meta
// 重置某些字段，例如模型名称等
if info.Request != nil {
info.Request.SetModelName(info.OriginModelName)
}
}
func (info *RelayInfo) ToString() string {
......
relay/relay-text.go → relay/compatible_handler.go
......@@ -25,38 +25,40 @@ import (
)
func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
textRequest, ok := info.Request.(*dto.GeneralOpenAIRequest)
textReq, ok := info.Request.(*dto.GeneralOpenAIRequest)
if !ok {
//return types.NewErrorWithStatusCode(errors.New("invalid request type"), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
common.FatalLog("invalid request type, expected dto.GeneralOpenAIRequest, got %T", info.Request)
return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected dto.GeneralOpenAIRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
request, err := common.DeepCopy(textReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to GeneralOpenAIRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
if textRequest.WebSearchOptions != nil {
c.Set("chat_completion_web_search_context_size", textRequest.WebSearchOptions.SearchContextSize)
if request.WebSearchOptions != nil {
c.Set("chat_completion_web_search_context_size", request.WebSearchOptions.SearchContextSize)
}
err := helper.ModelMappedHelper(c, info, textRequest)
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
includeUsage := true
// 判断用户是否需要返回使用情况
if textRequest.StreamOptions != nil {
includeUsage = textRequest.StreamOptions.IncludeUsage
if request.StreamOptions != nil {
includeUsage = request.StreamOptions.IncludeUsage
}
// 如果不支持StreamOptions，将StreamOptions设置为nil
if !info.SupportStreamOptions || !textRequest.Stream {
textRequest.StreamOptions = nil
if !info.SupportStreamOptions || !request.Stream {
request.StreamOptions = nil
} else {
// 如果支持StreamOptions，且请求中没有设置StreamOptions，根据配置文件设置StreamOptions
if constant.ForceStreamOption {
textRequest.StreamOptions = &dto.StreamOptions{
request.StreamOptions = &dto.StreamOptions{
IncludeUsage: true,
}
}
......@@ -81,7 +83,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
}
requestBody = bytes.NewBuffer(body)
} else {
convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, textRequest)
convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
......
relay/embedding_handler.go
......@@ -16,15 +16,19 @@ import (
)
func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
embeddingRequest, ok := info.Request.(*dto.EmbeddingRequest)
embeddingReq, ok := info.Request.(*dto.EmbeddingRequest)
if !ok {
common.FatalLog(fmt.Sprintf("invalid request type, expected *dto.EmbeddingRequest, got %T", info.Request))
return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected *dto.EmbeddingRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
request, err := common.DeepCopy(embeddingReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to EmbeddingRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
err := helper.ModelMappedHelper(c, info, embeddingRequest)
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
......@@ -35,7 +39,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
}
adaptor.Init(info)
convertedRequest, err := adaptor.ConvertEmbeddingRequest(c, info, *embeddingRequest)
convertedRequest, err := adaptor.ConvertEmbeddingRequest(c, info, *request)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
......
relay/gemini_handler.go
......@@ -53,13 +53,18 @@ func trimModelThinking(modelName string) string {
func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
request, ok := info.Request.(*dto.GeminiChatRequest)
geminiReq, ok := info.Request.(*dto.GeminiChatRequest)
if !ok {
common.FatalLog(fmt.Sprintf("invalid request type, expected *dto.GeminiChatRequest, got %T", info.Request))
return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected *dto.GeminiChatRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
request, err := common.DeepCopy(geminiReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to GeminiChatRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
// model mapped 模型映射
err := helper.ModelMappedHelper(c, info, request)
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
......@@ -170,7 +175,7 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
isBatch := strings.HasSuffix(c.Request.URL.Path, "batchEmbedContents")
info.IsGeminiBatchEmbedding = isBatch
var req any
var req dto.Request
var err error
var inputTexts []string
......
relay/helper/model_mapped.go
......@@ -4,15 +4,12 @@ import (
"encoding/json"
"errors"
"fmt"
"github.com/gin-gonic/gin"
"one-api/dto"
common2 "one-api/logger"
"one-api/relay/common"
"one-api/types"
"github.com/gin-gonic/gin"
)
func ModelMappedHelper(c *gin.Context, info *common.RelayInfo, request any) error {
func ModelMappedHelper(c *gin.Context, info *common.RelayInfo, request dto.Request) error {
// map model name
modelMapping := c.GetString("model_mapping")
if modelMapping != "" && modelMapping != "{}" {
......@@ -54,40 +51,7 @@ func ModelMappedHelper(c *gin.Context, info *common.RelayInfo, request any) erro
}
}
if request != nil {
switch info.RelayFormat {
case types.RelayFormatGemini:
// Gemini 模型映射
case types.RelayFormatClaude:
if claudeRequest, ok := request.(*dto.ClaudeRequest); ok {
claudeRequest.Model = info.UpstreamModelName
}
case types.RelayFormatOpenAIResponses:
if openAIResponsesRequest, ok := request.(*dto.OpenAIResponsesRequest); ok {
openAIResponsesRequest.Model = info.UpstreamModelName
}
case types.RelayFormatOpenAIAudio:
if openAIAudioRequest, ok := request.(*dto.AudioRequest); ok {
openAIAudioRequest.Model = info.UpstreamModelName
}
case types.RelayFormatOpenAIImage:
if imageRequest, ok := request.(*dto.ImageRequest); ok {
imageRequest.Model = info.UpstreamModelName
}
case types.RelayFormatRerank:
if rerankRequest, ok := request.(*dto.RerankRequest); ok {
rerankRequest.Model = info.UpstreamModelName
}
case types.RelayFormatEmbedding:
if embeddingRequest, ok := request.(*dto.EmbeddingRequest); ok {
embeddingRequest.Model = info.UpstreamModelName
}
default:
if openAIRequest, ok := request.(*dto.GeneralOpenAIRequest); ok {
openAIRequest.Model = info.UpstreamModelName
} else {
common2.LogWarn(c, fmt.Sprintf("model mapped but request type %T not supported", request))
}
}
request.SetModelName(info.UpstreamModelName)
}
return nil
}
relay/image_handler.go
......@@ -20,16 +20,19 @@ import (
)
func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
imageRequest, ok := info.Request.(*dto.ImageRequest)
imageReq, ok := info.Request.(*dto.ImageRequest)
if !ok {
common.FatalLog(fmt.Sprintf("invalid request type, expected dto.ImageRequest, got %T", info.Request))
return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected dto.ImageRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
request, err := common.DeepCopy(imageReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to ImageRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
err := helper.ModelMappedHelper(c, info, imageRequest)
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
......@@ -49,7 +52,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
}
requestBody = bytes.NewBuffer(body)
} else {
convertedRequest, err := adaptor.ConvertImageRequest(c, info, *imageRequest)
convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed)
}
......@@ -102,21 +105,21 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
}
if usage.(*dto.Usage).TotalTokens == 0 {
usage.(*dto.Usage).TotalTokens = int(imageRequest.N)
usage.(*dto.Usage).TotalTokens = int(request.N)
}
if usage.(*dto.Usage).PromptTokens == 0 {
usage.(*dto.Usage).PromptTokens = int(imageRequest.N)
usage.(*dto.Usage).PromptTokens = int(request.N)
}
quality := "standard"
if imageRequest.Quality == "hd" {
if request.Quality == "hd" {
quality = "hd"
}
var logContent string
if len(imageRequest.Size) > 0 {
logContent = fmt.Sprintf("大小 %s, 品质 %s", imageRequest.Size, quality)
if len(request.Size) > 0 {
logContent = fmt.Sprintf("大小 %s, 品质 %s", request.Size, quality)
}
postConsumeQuota(c, info, usage.(*dto.Usage), logContent)
......
relay/rerank_handler.go
......@@ -16,23 +16,20 @@ import (
"github.com/gin-gonic/gin"
)
func getRerankPromptToken(rerankRequest dto.RerankRequest) int {
token := service.CountTokenInput(rerankRequest.Query, rerankRequest.Model)
for _, document := range rerankRequest.Documents {
tkm := service.CountTokenInput(document, rerankRequest.Model)
token += tkm
}
return token
}
func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
rerankRequest, ok := info.Request.(*dto.RerankRequest)
rerankReq, ok := info.Request.(*dto.RerankRequest)
if !ok {
common.FatalLog(fmt.Sprintf("invalid request type, expected dto.RerankRequest, got %T", info.Request))
return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected dto.RerankRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
request, err := common.DeepCopy(rerankReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to ImageRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
err := helper.ModelMappedHelper(c, info, rerankRequest)
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
......@@ -51,7 +48,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
requestBody = bytes.NewBuffer(body)
} else {
convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *rerankRequest)
convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *request)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
......
relay/responses_handler.go
......@@ -20,12 +20,17 @@ import (
func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
request, ok := info.Request.(*dto.OpenAIResponsesRequest)
responsesReq, ok := info.Request.(*dto.OpenAIResponsesRequest)
if !ok {
common.FatalLog(fmt.Sprintf("invalid request type, expected dto.OpenAIResponsesRequest, got %T", info.Request))
return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected dto.OpenAIResponsesRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
err := helper.ModelMappedHelper(c, info, request)
request, err := common.DeepCopy(responsesReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to GeneralOpenAIRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment