refactor(antigravity): 简化模型映射逻辑,支持前缀匹配

- 删除精确映射表 antigravityModelMapping,统一使用前缀映射
- 前缀映射支持模型版本号变化(如 -20251111, -thinking, -preview)
- 简化 IsModelSupported 函数,所有 claude-/gemini- 前缀模型都支持
- 添加跨协议测试用例:Claude 端点调用 Gemini 模型、Gemini 端点调用 Claude 模型
This commit is contained in:
song
2025-12-31 21:16:32 +08:00
parent 312cc00d21
commit 7f7bbdf677
3 changed files with 96 additions and 59 deletions

View File

@@ -57,6 +57,7 @@ var geminiModels = []string{
"gemini-2.5-flash-lite", "gemini-2.5-flash-lite",
"gemini-3-flash", "gemini-3-flash",
"gemini-3-pro-low", "gemini-3-pro-low",
"gemini-3-pro-high",
} }
func TestMain(m *testing.M) { func TestMain(m *testing.M) {
@@ -641,6 +642,37 @@ func testClaudeThinkingWithToolHistory(t *testing.T, model string) {
t.Logf("✅ thinking 模式工具调用测试通过, id=%v", result["id"]) t.Logf("✅ thinking 模式工具调用测试通过, id=%v", result["id"])
} }
// TestClaudeMessagesWithGeminiModel 测试在 Claude 端点使用 Gemini 模型
// 验证:通过 /v1/messages 端点传入 gemini 模型名的场景(含前缀映射)
// 仅在 Antigravity 模式下运行ENDPOINT_PREFIX="/antigravity"
func TestClaudeMessagesWithGeminiModel(t *testing.T) {
if endpointPrefix != "/antigravity" {
t.Skip("仅在 Antigravity 模式下运行")
}
// 测试通过 Claude 端点调用 Gemini 模型
geminiViaClaude := []string{
"gemini-3-flash", // 直接支持
"gemini-3-pro-low", // 直接支持
"gemini-3-pro-high", // 直接支持
"gemini-3-pro", // 前缀映射 -> gemini-3-pro-high
"gemini-3-pro-preview", // 前缀映射 -> gemini-3-pro-high
}
for i, model := range geminiViaClaude {
if i > 0 {
time.Sleep(testInterval)
}
t.Run(model+"_通过Claude端点", func(t *testing.T) {
testClaudeMessage(t, model, false)
})
time.Sleep(testInterval)
t.Run(model+"_通过Claude端点_流式", func(t *testing.T) {
testClaudeMessage(t, model, true)
})
}
}
// TestClaudeMessagesWithNoSignature 测试历史 thinking block 不带 signature 的场景 // TestClaudeMessagesWithNoSignature 测试历史 thinking block 不带 signature 的场景
// 验证Gemini 模型接受没有 signature 的 thinking block // 验证Gemini 模型接受没有 signature 的 thinking block
func TestClaudeMessagesWithNoSignature(t *testing.T) { func TestClaudeMessagesWithNoSignature(t *testing.T) {
@@ -738,3 +770,30 @@ func testClaudeWithNoSignature(t *testing.T, model string) {
} }
t.Logf("✅ 无 signature thinking 处理测试通过, id=%v", result["id"]) t.Logf("✅ 无 signature thinking 处理测试通过, id=%v", result["id"])
} }
// TestGeminiEndpointWithClaudeModel 测试通过 Gemini 端点调用 Claude 模型
// 仅在 Antigravity 模式下运行ENDPOINT_PREFIX="/antigravity"
func TestGeminiEndpointWithClaudeModel(t *testing.T) {
if endpointPrefix != "/antigravity" {
t.Skip("仅在 Antigravity 模式下运行")
}
// 测试通过 Gemini 端点调用 Claude 模型
claudeViaGemini := []string{
"claude-sonnet-4-5",
"claude-opus-4-5-thinking",
}
for i, model := range claudeViaGemini {
if i > 0 {
time.Sleep(testInterval)
}
t.Run(model+"_通过Gemini端点", func(t *testing.T) {
testGeminiGenerate(t, model, false)
})
time.Sleep(testInterval)
t.Run(model+"_通过Gemini端点_流式", func(t *testing.T) {
testGeminiGenerate(t, model, true)
})
}
}

View File

@@ -25,7 +25,7 @@ const (
antigravityRetryMaxDelay = 16 * time.Second antigravityRetryMaxDelay = 16 * time.Second
) )
// Antigravity 直接支持的模型 // Antigravity 直接支持的模型(精确匹配透传)
var antigravitySupportedModels = map[string]bool{ var antigravitySupportedModels = map[string]bool{
"claude-opus-4-5-thinking": true, "claude-opus-4-5-thinking": true,
"claude-sonnet-4-5": true, "claude-sonnet-4-5": true,
@@ -36,23 +36,26 @@ var antigravitySupportedModels = map[string]bool{
"gemini-3-flash": true, "gemini-3-flash": true,
"gemini-3-pro-low": true, "gemini-3-pro-low": true,
"gemini-3-pro-high": true, "gemini-3-pro-high": true,
"gemini-3-pro-preview": true,
"gemini-3-pro-image": true, "gemini-3-pro-image": true,
} }
// Antigravity 系统默认模型映射表(不支持 → 支持 // Antigravity 前缀映射表(按前缀长度降序排列,确保最长匹配优先
var antigravityModelMapping = map[string]string{ // 用于处理模型版本号变化(如 -20251111, -thinking, -preview 等后缀)
"claude-3-5-sonnet-20241022": "claude-sonnet-4-5", var antigravityPrefixMapping = []struct {
"claude-3-5-sonnet-20240620": "claude-sonnet-4-5", prefix string
"claude-sonnet-4-5-20250929": "claude-sonnet-4-5-thinking", target string
"claude-opus-4": "claude-opus-4-5-thinking", }{
"claude-opus-4-5-20251101": "claude-opus-4-5-thinking", // 长前缀优先
"claude-haiku-4": "gemini-3-flash", {"gemini-3-pro-image", "gemini-3-pro-image"}, // gemini-3-pro-image-preview 等
"claude-haiku-4-5": "gemini-3-flash", {"claude-3-5-sonnet", "claude-sonnet-4-5"}, // 旧版 claude-3-5-sonnet-xxx
"claude-3-haiku-20240307": "gemini-3-flash", {"claude-sonnet-4-5", "claude-sonnet-4-5"}, // claude-sonnet-4-5-xxx
"claude-haiku-4-5-20251001": "gemini-3-flash", {"claude-haiku-4-5", "gemini-3-flash"}, // claude-haiku-4-5-xxx
// 生图模型:官方名 → Antigravity 内部名 {"claude-opus-4-5", "claude-opus-4-5-thinking"},
"gemini-3-pro-image-preview": "gemini-3-pro-image", {"claude-3-haiku", "gemini-3-flash"}, // 旧版 claude-3-haiku-xxx
{"claude-sonnet-4", "claude-sonnet-4-5"},
{"claude-haiku-4", "gemini-3-flash"},
{"claude-opus-4", "claude-opus-4-5-thinking"},
{"gemini-3-pro", "gemini-3-pro-high"}, // gemini-3-pro, gemini-3-pro-preview 等
} }
// AntigravityGatewayService 处理 Antigravity 平台的 API 转发 // AntigravityGatewayService 处理 Antigravity 平台的 API 转发
@@ -84,24 +87,27 @@ func (s *AntigravityGatewayService) GetTokenProvider() *AntigravityTokenProvider
} }
// getMappedModel 获取映射后的模型名 // getMappedModel 获取映射后的模型名
// 逻辑:账户映射 → 直接支持透传 → 前缀映射 → gemini透传 → 默认值
func (s *AntigravityGatewayService) getMappedModel(account *Account, requestedModel string) string { func (s *AntigravityGatewayService) getMappedModel(account *Account, requestedModel string) string {
// 1. 优先使用账户级映射(复用现有方法 // 1. 账户级映射(用户自定义优先
if mapped := account.GetMappedModel(requestedModel); mapped != requestedModel { if mapped := account.GetMappedModel(requestedModel); mapped != requestedModel {
return mapped return mapped
} }
// 2. 系统默认映射 // 2. 直接支持的模型透传
if mapped, ok := antigravityModelMapping[requestedModel]; ok { if antigravitySupportedModels[requestedModel] {
return mapped
}
// 3. Gemini 模型透传
if strings.HasPrefix(requestedModel, "gemini-") {
return requestedModel return requestedModel
} }
// 4. Claude 前缀透传直接支持的模型 // 3. 前缀映射(处理版本号变化,如 -20251111, -thinking, -preview
if antigravitySupportedModels[requestedModel] { for _, pm := range antigravityPrefixMapping {
if strings.HasPrefix(requestedModel, pm.prefix) {
return pm.target
}
}
// 4. Gemini 模型透传(未匹配到前缀的 gemini 模型)
if strings.HasPrefix(requestedModel, "gemini-") {
return requestedModel return requestedModel
} }
@@ -110,24 +116,10 @@ func (s *AntigravityGatewayService) getMappedModel(account *Account, requestedMo
} }
// IsModelSupported 检查模型是否被支持 // IsModelSupported 检查模型是否被支持
// 所有 claude- 和 gemini- 前缀的模型都能通过映射或透传支持
func (s *AntigravityGatewayService) IsModelSupported(requestedModel string) bool { func (s *AntigravityGatewayService) IsModelSupported(requestedModel string) bool {
// 直接支持的模型 return strings.HasPrefix(requestedModel, "claude-") ||
if antigravitySupportedModels[requestedModel] { strings.HasPrefix(requestedModel, "gemini-")
return true
}
// 可映射的模型
if _, ok := antigravityModelMapping[requestedModel]; ok {
return true
}
// Gemini 前缀透传
if strings.HasPrefix(requestedModel, "gemini-") {
return true
}
// Claude 模型支持(通过默认映射)
if strings.HasPrefix(requestedModel, "claude-") {
return true
}
return false
} }
// TestConnectionResult 测试连接结果 // TestConnectionResult 测试连接结果

View File

@@ -515,24 +515,10 @@ func (s *GatewayService) isModelSupportedByAccount(account *Account, requestedMo
} }
// IsAntigravityModelSupported 检查 Antigravity 平台是否支持指定模型 // IsAntigravityModelSupported 检查 Antigravity 平台是否支持指定模型
// 所有 claude- 和 gemini- 前缀的模型都能通过映射或透传支持
func IsAntigravityModelSupported(requestedModel string) bool { func IsAntigravityModelSupported(requestedModel string) bool {
// 直接支持的模型 return strings.HasPrefix(requestedModel, "claude-") ||
if antigravitySupportedModels[requestedModel] { strings.HasPrefix(requestedModel, "gemini-")
return true
}
// 可映射的模型
if _, ok := antigravityModelMapping[requestedModel]; ok {
return true
}
// Gemini 前缀透传
if strings.HasPrefix(requestedModel, "gemini-") {
return true
}
// Claude 模型支持(通过默认映射到 claude-sonnet-4-5
if strings.HasPrefix(requestedModel, "claude-") {
return true
}
return false
} }
// GetAccessToken 获取账号凭证 // GetAccessToken 获取账号凭证