feat: support model mapping chain

#1033
This commit is contained in:
joey
2025-05-07 16:00:35 +08:00
committed by GitHub
parent 65a15dbc17
commit 65ccfd0848

View File

@@ -2,9 +2,11 @@ package helper
import ( import (
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"github.com/gin-gonic/gin"
"one-api/relay/common" "one-api/relay/common"
"github.com/gin-gonic/gin"
) )
func ModelMappedHelper(c *gin.Context, info *common.RelayInfo) error { func ModelMappedHelper(c *gin.Context, info *common.RelayInfo) error {
@@ -16,9 +18,36 @@ func ModelMappedHelper(c *gin.Context, info *common.RelayInfo) error {
if err != nil { if err != nil {
return fmt.Errorf("unmarshal_model_mapping_failed") return fmt.Errorf("unmarshal_model_mapping_failed")
} }
if modelMap[info.OriginModelName] != "" {
info.UpstreamModelName = modelMap[info.OriginModelName] // 支持链式模型重定向,最终使用链尾的模型
info.IsModelMapped = true currentModel := info.OriginModelName
visitedModels := map[string]bool{
currentModel: true,
}
for {
if mappedModel, exists := modelMap[currentModel]; exists && mappedModel != "" {
// 模型重定向循环检测,避免无限循环
if visitedModels[mappedModel] {
if mappedModel == currentModel {
if currentModel == info.OriginModelName {
info.IsModelMapped = false
return nil
} else {
info.IsModelMapped = true
break
}
}
return errors.New("model_mapping_contains_cycle")
}
visitedModels[mappedModel] = true
currentModel = mappedModel
info.IsModelMapped = true
} else {
break
}
}
if info.IsModelMapped {
info.UpstreamModelName = currentModel
} }
} }
return nil return nil