fix(gemini): enable model_mapping filtering for Gemini API Key accounts
Remove the special case that bypassed model-supported checks for Gemini API Key accounts, allowing model_mapping to filter requests properly. Add tests for multiplatform model filtering behavior.
This commit is contained in:
@@ -895,6 +895,55 @@ func TestGatewayService_SelectAccountForModelWithPlatform_GeminiPreferOAuth(t *t
|
|||||||
require.Equal(t, int64(2), acc.ID)
|
require.Equal(t, int64(2), acc.ID)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestGatewayService_SelectAccountForModelWithPlatform_GeminiAPIKeyModelMappingFilter(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
repo := &mockAccountRepoForPlatform{
|
||||||
|
accounts: []Account{
|
||||||
|
{
|
||||||
|
ID: 1,
|
||||||
|
Platform: PlatformGemini,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Priority: 1,
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Credentials: map[string]any{"model_mapping": map[string]any{"gemini-2.5-pro": "gemini-2.5-pro"}},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
ID: 2,
|
||||||
|
Platform: PlatformGemini,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Priority: 2,
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Credentials: map[string]any{"model_mapping": map[string]any{"gemini-2.5-flash": "gemini-2.5-flash"}},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
accountsByID: map[int64]*Account{},
|
||||||
|
}
|
||||||
|
for i := range repo.accounts {
|
||||||
|
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||||
|
}
|
||||||
|
|
||||||
|
cache := &mockGatewayCacheForPlatform{}
|
||||||
|
|
||||||
|
svc := &GatewayService{
|
||||||
|
accountRepo: repo,
|
||||||
|
cache: cache,
|
||||||
|
cfg: testConfig(),
|
||||||
|
}
|
||||||
|
|
||||||
|
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "", "gemini-2.5-flash", nil, PlatformGemini)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, acc)
|
||||||
|
require.Equal(t, int64(2), acc.ID, "应过滤不支持请求模型的 APIKey 账号")
|
||||||
|
|
||||||
|
acc, err = svc.selectAccountForModelWithPlatform(ctx, nil, "", "gemini-3-pro-preview", nil, PlatformGemini)
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Nil(t, acc)
|
||||||
|
require.Contains(t, err.Error(), "supporting model")
|
||||||
|
}
|
||||||
|
|
||||||
func TestGatewayService_SelectAccountForModelWithPlatform_StickyInGroup(t *testing.T) {
|
func TestGatewayService_SelectAccountForModelWithPlatform_StickyInGroup(t *testing.T) {
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
groupID := int64(50)
|
groupID := int64(50)
|
||||||
@@ -1070,6 +1119,36 @@ func TestGatewayService_isModelSupportedByAccount(t *testing.T) {
|
|||||||
model: "claude-3-5-sonnet-20241022",
|
model: "claude-3-5-sonnet-20241022",
|
||||||
expected: true,
|
expected: true,
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
name: "Gemini平台-无映射配置-支持所有模型",
|
||||||
|
account: &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey},
|
||||||
|
model: "gemini-2.5-flash",
|
||||||
|
expected: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Gemini平台-有映射配置-只支持配置的模型",
|
||||||
|
account: &Account{
|
||||||
|
Platform: PlatformGemini,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Credentials: map[string]any{
|
||||||
|
"model_mapping": map[string]any{"gemini-2.5-pro": "gemini-2.5-pro"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
model: "gemini-2.5-flash",
|
||||||
|
expected: false,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Gemini平台-有映射配置-支持配置的模型",
|
||||||
|
account: &Account{
|
||||||
|
Platform: PlatformGemini,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Credentials: map[string]any{
|
||||||
|
"model_mapping": map[string]any{"gemini-2.5-pro": "gemini-2.5-pro"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
model: "gemini-2.5-pro",
|
||||||
|
expected: true,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, tt := range tests {
|
for _, tt := range tests {
|
||||||
|
|||||||
@@ -2825,10 +2825,6 @@ func (s *GatewayService) isModelSupportedByAccount(account *Account, requestedMo
|
|||||||
if account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey {
|
if account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey {
|
||||||
requestedModel = claude.NormalizeModelID(requestedModel)
|
requestedModel = claude.NormalizeModelID(requestedModel)
|
||||||
}
|
}
|
||||||
// Gemini API Key 账户直接透传,由上游判断模型是否支持
|
|
||||||
if account.Platform == PlatformGemini && account.Type == AccountTypeAPIKey {
|
|
||||||
return true
|
|
||||||
}
|
|
||||||
// 其他平台使用账户的模型支持检查
|
// 其他平台使用账户的模型支持检查
|
||||||
return account.IsModelSupported(requestedModel)
|
return account.IsModelSupported(requestedModel)
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user