middleware/distributor.go
Rama main · réplica de lectura del repositorio original.
package middleware
import (
"errors"
"fmt"
"io"
"net/http"
"slices"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
taskdto "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/MoeForge/moego/ojson"
"github.com/gin-gonic/gin"
)
type ModelRequest struct {
Model string `json:"model"`
Group string `json:"group,omitempty"`
}
func Distribute() func(c *gin.Context) {
return func(c *gin.Context) {
var channel *model.Channel
defer func() {
if c.Writer.Status() >= 400 {
service.RecordRequestPolicyTermination(c, types.NewErrorWithStatusCode(errors.New("request rejected"), types.ErrorCodeInvalidRequest, c.Writer.Status(), types.ErrOptionWithSkipRetry()))
}
}()
constraints := service.GetChannelConstraints(c)
constraints.AddFilter(taskdto.ChannelFilter{
Kind: taskdto.FilterRequestPath,
RequestPath: c.Request.URL.Path,
})
service.AppendTaskPluginIdentityFilter(c, c.GetString("expected_task_plugin_key"))
modelRequest, shouldSelectChannel, err := getModelRequest(c)
if err != nil {
abortWithOpenAiMessage(c, invalidRequestStatus(err), i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
return
}
_, pinned, _ := constraints.ResolvedPin()
if !pinned {
// Select a channel for the user
// check token model mapping
modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled)
if modelLimitEnable {
s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit)
if !ok {
// token model limit is empty, all models are not allowed
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess))
return
}
var tokenModelLimit map[string]bool
tokenModelLimit, ok = s.(map[string]bool)
if !ok {
tokenModelLimit = map[string]bool{}
}
if !TokenModelLimitAllows(tokenModelLimit, modelRequest.Model) {
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model}))
return
}
}
if shouldSelectChannel {
if modelRequest.Model == "" {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorModelNameRequired))
return
}
// check path is /pg/chat/completions
if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
playgroundRequest := &dto.PlayGroundRequest{}
err = common.UnmarshalBodyReusable(c, playgroundRequest)
if err != nil {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidPlayground, map[string]any{"Error": err.Error()}))
return
}
if playgroundRequest.Group != "" {
if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.Group != usingGroup {
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorGroupAccessDenied))
return
}
common.SetContextKey(c, constant.ContextKeyUsingGroup, playgroundRequest.Group)
}
}
}
}
if pinned || shouldSelectChannel {
usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
var selectErr *service.ChannelSelectError
if value, exists := c.Get(contextKeyTaskPluginEndpointSelection); exists && !pinned {
if selected, ok := value.(taskPluginEndpointSelection); ok && selected.Model == modelRequest.Model && selected.Channel != nil && selected.Channel.Status == common.ChannelStatusEnabled {
if matches, _ := model.ChannelSatisfiesFilters(selected.Channel, modelRequest.Model, constraints.Filters); matches {
channel = selected.Channel
}
}
}
if channel == nil {
channel, _, selectErr = service.SelectChannelForRequest(c, modelRequest.Model, &service.RetryParam{
Ctx: c,
ModelName: modelRequest.Model,
TokenGroup: usingGroup,
RequestPath: c.Request.URL.Path,
Retry: common.GetPointer(0),
})
}
if selectErr != nil {
if selectErr.FilterKind == taskdto.FilterTaskPluginIdentity {
logTaskPluginChannelDecision(c, selectErr.Channel, modelRequest.Model, "channel_rejected", "identity_mismatch")
}
message := selectErr.Message
if selectErr.NoAvailableChannel {
message = noAvailableChannelMessage(c, usingGroup, modelRequest.Model)
} else if selectErr.MessageID != "" {
message = i18n.T(c, selectErr.MessageID, selectErr.Params)
}
abortWithOpenAiMessage(c, selectErr.StatusCode, message, selectErr.Code)
return
}
}
common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
SetupContextForSelectedChannel(c, channel, modelRequest.Model)
c.Next()
if channel != nil && c.Writer != nil && c.Writer.Status() < http.StatusBadRequest {
service.RecordChannelAffinity(c, channel.Id)
}
}
}
// noAvailableChannelMessage explains a 503 for a task-plugin-claimed model.
// The response tells the caller the model is plugin-claimed without naming the
// plugin; the candidate plugin keys go to the server log under the request id.
func noAvailableChannelMessage(c *gin.Context, group, modelName string) string {
value, exists := c.Get(jsplugin.ContextKeyPinnedPlugin)
pinned, ok := value.(jsplugin.PinnedPlugin)
if exists && ok && pinned.Plugin != nil {
keys := []string{pinned.Plugin.Meta.Key}
if _, candidates := pinnedCandidatePlugins(c); len(candidates) > 0 {
keys = nil
for _, candidate := range candidates {
if candidate != nil {
keys = append(keys, candidate.Meta.Key)
}
}
}
logger.LogWarn(c, common.LogText("task_plugin subsystem=distribution event=no_available_channel group=%q model=%q plugins=%q reason=no_eligible_channel", group, modelName, strings.Join(keys, ",")))
return i18n.T(c, i18n.MsgDistributorNoAvailableChannelTaskPlugin, map[string]any{"Group": group, "Model": modelName})
}
return i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": group, "Model": modelName})
}
func channelMatchesExpectedTaskPlugin(c *gin.Context, channel *model.Channel, expected string) bool {
if channel == nil {
return false
}
if c != nil {
if _, matched := pinnedEndpointCandidateForChannel(c, channel, expected); matched {
return true
}
}
if expected == "" {
return channel.Type != constant.ChannelTypeTaskPlugin
}
if channel.Type == constant.ChannelTypeTaskPlugin || channel.Type == constant.ChannelTypeNewAPI {
return channel.GetSetting().BindsTaskPlugin(expected)
}
if c == nil {
return false
}
value, exists := c.Get(jsplugin.ContextKeyPinnedPlugin)
pinned, ok := value.(jsplugin.PinnedPlugin)
if !exists || !ok || pinned.Generation == nil || pinned.Plugin == nil || pinned.Plugin.Meta.Key != expected {
return false
}
plugin, ok := pinned.Generation.GetByChannelType(channel.Type)
return ok && plugin == pinned.Plugin
}
// pinnedCandidatePlugins returns the plugins a pinned shared host endpoint, or
// a native route several plugins declare, may execute, in ascending key order.
func pinnedCandidatePlugins(c *gin.Context) (*jsplugin.RoutingGeneration, []*jsplugin.LoadedPlugin) {
if value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint); exists {
pinned, ok := value.(jsplugin.PinnedEndpoint)
if !ok || pinned.Generation == nil || pinned.Plugin == nil {
return nil, nil
}
if len(pinned.Candidates) == 0 {
return pinned.Generation, []*jsplugin.LoadedPlugin{pinned.Plugin}
}
plugins := make([]*jsplugin.LoadedPlugin, 0, len(pinned.Candidates))
for _, candidate := range pinned.Candidates {
plugins = append(plugins, candidate.Plugin)
}
return pinned.Generation, plugins
}
if value, exists := c.Get(jsplugin.ContextKeyPinnedRoute); exists {
pinned, ok := value.(jsplugin.PinnedRoute)
if !ok || pinned.Generation == nil || len(pinned.Candidates) < 2 {
return nil, nil
}
plugins := make([]*jsplugin.LoadedPlugin, 0, len(pinned.Candidates))
for _, candidate := range pinned.Candidates {
plugins = append(plugins, candidate.Plugin)
}
return pinned.Generation, plugins
}
return nil, nil
}
// pinnedEndpointCandidateForChannel returns the index, among
// pinnedCandidatePlugins, of the candidate that executes on channel.
func pinnedEndpointCandidateForChannel(c *gin.Context, channel *model.Channel, expected string) (int, bool) {
if c == nil || channel == nil || expected == "" {
return -1, false
}
generation, candidates := pinnedCandidatePlugins(c)
if generation == nil {
return -1, false
}
expectedOwned := false
selected := -1
setting := channel.GetSetting()
for index, candidate := range candidates {
if candidate == nil {
continue
}
if candidate.Meta.Key == expected {
expectedOwned = true
}
if channel.Type == constant.ChannelTypeTaskPlugin || channel.Type == constant.ChannelTypeNewAPI {
// A New API channel may bind several candidates. The first bound
// candidate in generation order executes, so the billing provider
// depends only on the channel and the request, never on which
// channels an earlier retry attempt happened to try.
if selected < 0 && setting.BindsTaskPlugin(candidate.Meta.Key) {
selected = index
}
continue
}
plugin, indexed := generation.GetByChannelType(channel.Type)
if indexed && plugin == candidate {
selected = index
}
}
return selected, expectedOwned && selected >= 0
}
// getModelFromRequest 从请求中读取模型信息
// 根据 Content-Type 自动处理:
// - application/json
// - application/x-www-form-urlencoded
// - multipart/form-data
func getModelFromRequest(c *gin.Context) (*ModelRequest, error) {
if cached, exists := c.Get(contextKeyTaskPluginEndpointModel); exists {
if modelRequest, ok := cached.(ModelRequest); ok {
cachedRequest := modelRequest
return &cachedRequest, nil
}
}
if strings.HasPrefix(c.Request.Header.Get("Content-Type"), "application/json") {
return getModelFromJSONBody(c)
}
var modelRequest ModelRequest
err := common.UnmarshalBodyReusable(c, &modelRequest)
if err != nil {
return nil, err
}
return &modelRequest, nil
}
func getModelFromJSONBody(c *gin.Context) (*ModelRequest, error) {
storage, err := common.GetBodyStorage(c)
if err != nil {
return nil, err
}
requestBody, err := storage.Bytes()
if err != nil {
return nil, err
}
if !ojson.ValidBytes(requestBody) {
return nil, errors.New("invalid JSON request body")
}
if countTopLevelJSONKey(requestBody, "model") > 1 {
return nil, errors.New("model must be provided once")
}
values := ojson.GetManyBytes(requestBody, "model", "group")
model, err := getJSONStringValue(values[0], "model")
if err != nil {
return nil, err
}
group, err := getJSONStringValue(values[1], "group")
if err != nil {
return nil, err
}
if _, seekErr := storage.Seek(0, io.SeekStart); seekErr != nil {
return nil, seekErr
}
c.Request.Body = io.NopCloser(storage)
return &ModelRequest{
Model: model,
Group: group,
}, nil
}
func countTopLevelJSONKey(data []byte, target string) int {
depth := 0
inString := false
escaped := false
stringStart := 0
expectingKey := false
count := 0
for index, current := range data {
if inString {
if escaped {
escaped = false
continue
}
if current == '\\' {
escaped = true
continue
}
if current != '"' {
continue
}
inString = false
if depth == 1 && expectingKey {
key := string(data[stringStart:index])
var decodedKey string
if common.Unmarshal(data[stringStart-1:index+1], &decodedKey) == nil {
key = decodedKey
}
cursor := index + 1
for cursor < len(data) && (data[cursor] == ' ' || data[cursor] == '\t' || data[cursor] == '\r' || data[cursor] == '\n') {
cursor++
}
if cursor < len(data) && data[cursor] == ':' && key == target {
count++
}
expectingKey = false
}
continue
}
switch current {
case '"':
inString = true
stringStart = index + 1
case '{':
depth++
if depth == 1 {
expectingKey = true
}
case '}':
depth--
case ',':
if depth == 1 {
expectingKey = true
}
}
}
return count
}
func getJSONStringValue(result ojson.Result, field string) (string, error) {
if !result.Exists() || result.Type == ojson.Null {
return "", nil
}
if result.Type != ojson.String {
return "", fmt.Errorf("field %s must be a string", field)
}
return result.String(), nil
}
func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
var modelRequest ModelRequest
shouldSelectChannel := true
var err error
if modelName := c.GetString("resolved_task_model"); modelName != "" {
modelRequest.Model = modelName
} else if strings.Contains(c.Request.URL.Path, "/mj/") {
relayMode := relayconstant.Path2RelayModeMidjourney(c.Request.URL.Path)
if relayMode == relayconstant.RelayModeMidjourneyTaskFetch ||
relayMode == relayconstant.RelayModeMidjourneyTaskFetchByCondition ||
relayMode == relayconstant.RelayModeMidjourneyNotify ||
relayMode == relayconstant.RelayModeMidjourneyTaskImageSeed {
shouldSelectChannel = false
} else {
midjourneyRequest := taskdto.MidjourneyRequest{}
err = common.UnmarshalBodyReusable(c, &midjourneyRequest)
if err != nil {
return nil, false, err
}
midjourneyModel, mjErr, success := service.GetMjRequestModel(relayMode, &midjourneyRequest)
if mjErr != nil {
return nil, false, fmt.Errorf("%s", mjErr.Description)
}
if midjourneyModel == "" {
if !success {
return nil, false, fmt.Errorf("%s", i18n.T(c, i18n.MsgDistributorInvalidParseModel))
} else {
// task fetch, task fetch by condition, notify
shouldSelectChannel = false
}
}
modelRequest.Model = midjourneyModel
}
c.Set("relay_mode", relayMode)
} else if strings.Contains(c.Request.URL.Path, "/v1/videos/") && strings.HasSuffix(c.Request.URL.Path, "/remix") {
relayMode := relayconstant.RelayModeVideoSubmit
c.Set("relay_mode", relayMode)
shouldSelectChannel = false
} else if strings.Contains(c.Request.URL.Path, "/v1/videos") {
//curl https://api.openai.com/v1/videos \
// -H "Authorization: Bearer $OPENAI_API_KEY" \
// -F "model=sora-2" \
// -F "prompt=A calico cat playing a piano on stage"
// -F input_reference="@image.jpg"
relayMode := relayconstant.RelayModeUnknown
if c.Request.Method == http.MethodPost {
relayMode = relayconstant.RelayModeVideoSubmit
req, err := getModelFromRequest(c)
if err != nil {
return nil, false, err
}
if req != nil {
modelRequest.Model = req.Model
}
} else if c.Request.Method == http.MethodGet {
relayMode = relayconstant.RelayModeVideoFetchByID
shouldSelectChannel = false
modelRequest.Model = getTaskOriginModelName(c)
}
c.Set("relay_mode", relayMode)
} else if strings.Contains(c.Request.URL.Path, "/v1/video/generations") {
relayMode := relayconstant.RelayModeUnknown
if c.Request.Method == http.MethodPost {
req, err := getModelFromRequest(c)
if err != nil {
return nil, false, err
}
modelRequest.Model = req.Model
relayMode = relayconstant.RelayModeVideoSubmit
} else if c.Request.Method == http.MethodGet {
relayMode = relayconstant.RelayModeVideoFetchByID
shouldSelectChannel = false
modelRequest.Model = getTaskOriginModelName(c)
}
if _, ok := c.Get("relay_mode"); !ok {
c.Set("relay_mode", relayMode)
}
} else if strings.HasPrefix(c.Request.URL.Path, "/v1beta/models/") || strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
// Gemini API 路径处理: /v1beta/models/gemini-2.0-flash:generateContent
relayMode := relayconstant.RelayModeGemini
modelName := extractModelNameFromGeminiPath(c.Request.URL.Path)
if modelName != "" {
modelRequest.Model = modelName
}
c.Set("relay_mode", relayMode)
} else if !strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") && !strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
req, err := getModelFromRequest(c)
if err != nil {
return nil, false, err
}
modelRequest.Model = req.Model
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/realtime") {
//wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01
modelRequest.Model = c.Query("model")
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") {
if modelRequest.Model == "" {
modelRequest.Model = "text-moderation-stable"
}
}
if strings.HasSuffix(c.Request.URL.Path, "embeddings") {
if modelRequest.Model == "" {
modelRequest.Model = c.Param("model")
}
}
// The routing name is also the billing name, so a model named in the body
// must not replace it: DeepSeek ignores that field.
if strings.HasPrefix(c.Request.URL.Path, "/v1/tools/search") {
modelRequest.Model = constant.DeepSeekSearchModel
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/tools/fetch") {
modelRequest.Model = constant.DeepSeekFetchModel
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") {
modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "dall-e")
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/images/edits") {
//modelRequest.Model = common.GetStringIfEmpty(c.PostForm("model"), "gpt-image-1")
contentType := c.ContentType()
if slices.Contains([]string{gin.MIMEPOSTForm, gin.MIMEMultipartPOSTForm}, contentType) {
req, err := getModelFromRequest(c)
if err == nil && req.Model != "" {
modelRequest.Model = req.Model
}
}
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/audio") {
relayMode := relayconstant.RelayModeAudioSpeech
if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/speech") {
modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "tts-1")
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") {
// 先尝试从请求读取
if req, err := getModelFromRequest(c); err == nil && req.Model != "" {
modelRequest.Model = req.Model
}
modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "whisper-1")
relayMode = relayconstant.RelayModeAudioTranslation
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") {
// 先尝试从请求读取
if req, err := getModelFromRequest(c); err == nil && req.Model != "" {
modelRequest.Model = req.Model
}
modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "whisper-1")
relayMode = relayconstant.RelayModeAudioTranscription
}
c.Set("relay_mode", relayMode)
}
if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
// playground chat completions
req, err := getModelFromRequest(c)
if err != nil {
return nil, false, err
}
modelRequest.Model = req.Model
modelRequest.Group = req.Group
common.SetContextKey(c, constant.ContextKeyTokenGroup, modelRequest.Group)
}
return &modelRequest, shouldSelectChannel, nil
}
// TokenModelLimitAllows reports whether a token model-limit map authorizes
// model. Exact name, wildcard-normalized name, and routing-normalized name
// (modifiers and legacy aliases stripped) are all accepted. The Responses
// WebSocket relay shares this rule so both transports admit the same names.
func TokenModelLimitAllows(limit map[string]bool, model string) bool {
if limit[model] {
return true
}
if formatted := ratio_setting.FormatMatchingModelName(model); limit[formatted] {
return true
}
return limit[ratio_setting.RoutingMatchModelName(model)]
}
// 修复 #4834: GET /v1/video/generations/:task_id && /v1/video/:task_id 此前不解析 model,
// 当 token 启用「可用模型限制」时,下游 modelLimitEnable 校验会因
// modelRequest.Model 为空而误报 "This token has no access to model"。
// 从已存储的任务记录中回填 OriginModelName 即可让校验走在正确的模型上。
func getTaskOriginModelName(c *gin.Context) string {
if !common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) {
return ""
}
taskId := c.Param("task_id")
if taskId == "" {
return ""
}
userId := c.GetInt("id")
if task, exist, err := model.GetByTaskId(userId, taskId); err == nil && exist && task != nil {
return task.Properties.OriginModelName
}
return ""
}
func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError {
c.Set("original_model", modelName) // for retry
expectedPlugin := c.GetString("expected_task_plugin_key")
if channel == nil {
logTaskPluginChannelDecision(c, nil, modelName, "channel_rejected", "nil_channel")
return types.NewError(errors.New("channel is nil"), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
}
if expectedPlugin != "" && !channelMatchesExpectedTaskPlugin(c, channel, expectedPlugin) {
logTaskPluginChannelDecision(c, channel, modelName, "channel_rejected", "identity_mismatch")
return types.NewError(
errors.New("selected channel does not match the pinned task plugin"),
types.ErrorCodeGetChannelFailed,
types.ErrOptionWithSkipRetry(),
)
}
if index, matched := pinnedEndpointCandidateForChannel(c, channel, expectedPlugin); matched {
previousPlugin := ""
var selected *jsplugin.LoadedPlugin
var generation *jsplugin.RoutingGeneration
if value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint); exists {
if pinned, ok := value.(jsplugin.PinnedEndpoint); ok && index < len(pinned.Candidates) && pinned.Candidates[index].Plugin != pinned.Plugin {
candidate := pinned.Candidates[index]
previousPlugin, selected, generation = pinned.Plugin.Meta.Key, candidate.Plugin, pinned.Generation
pinned.Plugin = candidate.Plugin
pinned.Protocol = candidate.Protocol
pinned.Operation = candidate.Operation
c.Set(jsplugin.ContextKeyPinnedEndpoint, pinned)
}
} else if value, exists := c.Get(jsplugin.ContextKeyPinnedRoute); exists {
if pinned, ok := value.(jsplugin.PinnedRoute); ok && pinned.Plugin != nil && pinned.Candidates[index].Plugin != pinned.Plugin {
candidate := pinned.Candidates[index]
previousPlugin, selected, generation = pinned.Plugin.Meta.Key, candidate.Plugin, pinned.Generation
pinned.Plugin = candidate.Plugin
pinned.Route = candidate.Route
c.Set(jsplugin.ContextKeyPinnedRoute, pinned)
}
}
if selected != nil {
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: generation, Plugin: selected})
c.Set("expected_task_plugin_key", selected.Meta.Key)
c.Set("task_plugin_key", selected.Meta.Key)
c.Set("platform", selected.Meta.Key)
logger.LogDebug(
c,
"task_plugin subsystem=endpoint event=provider_selected generation=%d previous_plugin=%q plugin=%q model=%q channel_id=%d channel_type=%d",
generation.Number,
previousPlugin,
selected.Meta.Key,
modelName,
channel.Id,
channel.Type,
)
}
}
common.SetContextKey(c, constant.ContextKeyChannelId, channel.Id)
common.SetContextKey(c, constant.ContextKeyChannelName, channel.Name)
common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type)
common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime)
common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting())
common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings())
switch channel.Type {
case constant.ChannelTypeTaskPlugin:
c.Set("task_plugin_key", channel.GetSetting().TaskPluginKey)
case constant.ChannelTypeNewAPI:
// The bound plugin verified above executes; ordinary requests through
// the same gateway channel carry no pinned plugin.
if executing := c.GetString("expected_task_plugin_key"); executing != "" {
c.Set("task_plugin_key", executing)
}
}
logTaskPluginChannelDecision(c, channel, modelName, "channel_selected", "")
paramOverride := channel.GetParamOverride()
headerOverride := channel.GetHeaderOverride()
if mergedParam, applied := service.ApplyChannelAffinityOverrideTemplate(c, paramOverride); applied {
paramOverride = mergedParam
}
common.SetContextKey(c, constant.ContextKeyChannelParamOverride, paramOverride)
common.SetContextKey(c, constant.ContextKeyChannelHeaderOverride, headerOverride)
if nil != channel.OpenAIOrganization && *channel.OpenAIOrganization != "" {
common.SetContextKey(c, constant.ContextKeyChannelOrganization, *channel.OpenAIOrganization)
}
common.SetContextKey(c, constant.ContextKeyChannelAutoBan, channel.GetAutoBan())
common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping())
common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
key, index, newAPIError := channel.GetNextEnabledKey()
if newAPIError != nil {
return newAPIError
}
if channel.ChannelInfo.IsMultiKey {
common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true)
common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, index)
} else {
// 必须设置为 false,否则在重试到单个 key 的时候会导致日志显示错误
common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, false)
}
// c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key))
common.SetContextKey(c, constant.ContextKeyChannelKey, key)
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL())
common.SetContextKey(c, constant.ContextKeySystemPromptOverride, false)
// TODO: api_version统一
switch channel.Type {
case constant.ChannelTypeAzure:
c.Set("api_version", channel.Other)
case constant.ChannelTypeVertexAi:
c.Set("region", channel.Other)
case constant.ChannelTypeXunfei:
c.Set("api_version", channel.Other)
case constant.ChannelTypeGemini:
c.Set("api_version", channel.Other)
case constant.ChannelTypeAli:
c.Set("plugin", channel.Other)
case constant.ChannelCloudflare:
c.Set("api_version", channel.Other)
case constant.ChannelTypeMokaAI:
c.Set("api_version", channel.Other)
case constant.ChannelTypeCoze:
c.Set("bot_id", channel.Other)
}
return nil
}
// extractModelNameFromGeminiPath 从 Gemini API URL 路径中提取模型名
// 输入格式: /v1beta/models/gemini-2.0-flash:generateContent
// 输出: gemini-2.0-flash
func extractModelNameFromGeminiPath(path string) string {
// 查找 "/models/" 的位置
modelsPrefix := "/models/"
modelsIndex := strings.Index(path, modelsPrefix)
if modelsIndex == -1 {
return ""
}
// 从 "/models/" 之后开始提取
startIndex := modelsIndex + len(modelsPrefix)
if startIndex >= len(path) {
return ""
}
// 查找 ":" 的位置,模型名在 ":" 之前
colonIndex := strings.Index(path[startIndex:], ":")
if colonIndex == -1 {
// 如果没有找到 ":",返回从 "/models/" 到路径结尾的部分
return path[startIndex:]
}
// 返回模型名部分
return path[startIndex : startIndex+colonIndex]
}