fix: 修复 Gemini 风格 /v1/models 列表请求 (#6199)
* fix: support Gemini model listing on v1 route * test: cover Gemini query key model listing
This commit is contained in:
+2
-1
@@ -374,7 +374,8 @@ func TokenAuth() func(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
// gemini api 从query中获取key
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/v1beta/models") ||
|
||||
if c.Request.URL.Path == "/v1/models" ||
|
||||
strings.HasPrefix(c.Request.URL.Path, "/v1beta/models") ||
|
||||
strings.HasPrefix(c.Request.URL.Path, "/v1beta/openai/models") ||
|
||||
strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
|
||||
skKey := c.Query("key")
|
||||
|
||||
@@ -25,7 +25,7 @@ func SetRelayRouter(router *gin.Engine) {
|
||||
case c.GetHeader("x-api-key") != "" && c.GetHeader("anthropic-version") != "":
|
||||
controller.ListModels(c, constant.ChannelTypeAnthropic)
|
||||
case c.GetHeader("x-goog-api-key") != "" || c.Query("key") != "": // 单独的适配
|
||||
controller.RetrieveModel(c, constant.ChannelTypeGemini)
|
||||
controller.ListModels(c, constant.ChannelTypeGemini)
|
||||
default:
|
||||
controller.ListModels(c, constant.ChannelTypeOpenAI)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestListModelsSupportsOpenAIAndGeminiAuthentication(t *testing.T) {
|
||||
setupRelayRouterTestDB(t)
|
||||
|
||||
user := model.User{
|
||||
Username: "models-user",
|
||||
Status: common.UserStatusEnabled,
|
||||
Group: "default",
|
||||
Quota: 100,
|
||||
}
|
||||
require.NoError(t, model.DB.Create(&user).Error)
|
||||
require.NoError(t, model.DB.Create(&model.Token{
|
||||
UserId: user.Id,
|
||||
Key: "modelstestkey",
|
||||
Status: common.TokenStatusEnabled,
|
||||
ExpiredTime: -1,
|
||||
UnlimitedQuota: true,
|
||||
}).Error)
|
||||
|
||||
engine := gin.New()
|
||||
SetRelayRouter(engine)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
headerName string
|
||||
expectedObject string
|
||||
expectedField string
|
||||
}{
|
||||
{
|
||||
name: "OpenAI bearer token",
|
||||
path: "/v1/models",
|
||||
headerName: "Authorization",
|
||||
expectedObject: "list",
|
||||
expectedField: "data",
|
||||
},
|
||||
{
|
||||
name: "Gemini API key header",
|
||||
path: "/v1/models",
|
||||
headerName: "x-goog-api-key",
|
||||
expectedField: "models",
|
||||
},
|
||||
{
|
||||
name: "Gemini API key query",
|
||||
path: "/v1/models?key=modelstestkey",
|
||||
expectedField: "models",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodGet, test.path, nil)
|
||||
if test.headerName != "" {
|
||||
value := "modelstestkey"
|
||||
if test.headerName == "Authorization" {
|
||||
value = "Bearer " + value
|
||||
}
|
||||
request.Header.Set(test.headerName, value)
|
||||
}
|
||||
|
||||
engine.ServeHTTP(recorder, request)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var payload map[string]any
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &payload))
|
||||
assert.Contains(t, payload, test.expectedField)
|
||||
assert.NotContains(t, payload, "error")
|
||||
if test.expectedObject != "" {
|
||||
assert.Equal(t, test.expectedObject, payload["object"])
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func setupRelayRouterTestDB(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
originalIsMasterNode := common.IsMasterNode
|
||||
originalRedisEnabled := common.RedisEnabled
|
||||
originalSQLitePath := common.SQLitePath
|
||||
originalMainDatabaseType := common.MainDatabaseType()
|
||||
originalLogDatabaseType := common.LogDatabaseType()
|
||||
originalSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")
|
||||
|
||||
common.IsMasterNode = false
|
||||
common.RedisEnabled = false
|
||||
common.SQLitePath = fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
|
||||
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
||||
require.NoError(t, os.Setenv("SQL_DSN", "local"))
|
||||
require.NoError(t, model.InitDB())
|
||||
model.LOG_DB = model.DB
|
||||
require.NoError(t, model.DB.AutoMigrate(&model.User{}, &model.Token{}, &model.Ability{}))
|
||||
|
||||
t.Cleanup(func() {
|
||||
if sqlDB, err := model.DB.DB(); err == nil {
|
||||
_ = sqlDB.Close()
|
||||
}
|
||||
common.IsMasterNode = originalIsMasterNode
|
||||
common.RedisEnabled = originalRedisEnabled
|
||||
common.SQLitePath = originalSQLitePath
|
||||
common.SetDatabaseTypes(originalMainDatabaseType, originalLogDatabaseType)
|
||||
if hadSQLDSN {
|
||||
require.NoError(t, os.Setenv("SQL_DSN", originalSQLDSN))
|
||||
} else {
|
||||
require.NoError(t, os.Unsetenv("SQL_DSN"))
|
||||
}
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user