fix(gemini): add customtools fallback metadata

This commit is contained in:
YanzheL
2026-04-01 02:19:10 +08:00
parent 83a16dec19
commit 095bef9554
2 changed files with 35 additions and 1 deletions

View File

@@ -2,6 +2,8 @@
// It is used when upstream model listing is unavailable (e.g. OAuth token missing AI Studio scopes). // It is used when upstream model listing is unavailable (e.g. OAuth token missing AI Studio scopes).
package gemini package gemini
import "strings"
type Model struct { type Model struct {
Name string `json:"name"` Name string `json:"name"`
DisplayName string `json:"displayName,omitempty"` DisplayName string `json:"displayName,omitempty"`
@@ -23,10 +25,27 @@ func DefaultModels() []Model {
{Name: "models/gemini-3-flash-preview", SupportedGenerationMethods: methods}, {Name: "models/gemini-3-flash-preview", SupportedGenerationMethods: methods},
{Name: "models/gemini-3-pro-preview", SupportedGenerationMethods: methods}, {Name: "models/gemini-3-pro-preview", SupportedGenerationMethods: methods},
{Name: "models/gemini-3.1-pro-preview", SupportedGenerationMethods: methods}, {Name: "models/gemini-3.1-pro-preview", SupportedGenerationMethods: methods},
{Name: "models/gemini-3.1-pro-preview-customtools", SupportedGenerationMethods: methods},
{Name: "models/gemini-3.1-flash-image", SupportedGenerationMethods: methods}, {Name: "models/gemini-3.1-flash-image", SupportedGenerationMethods: methods},
} }
} }
func HasFallbackModel(model string) bool {
trimmed := strings.TrimSpace(model)
if trimmed == "" {
return false
}
if !strings.HasPrefix(trimmed, "models/") {
trimmed = "models/" + trimmed
}
for _, model := range DefaultModels() {
if model.Name == trimmed {
return true
}
}
return false
}
func FallbackModelsList() ModelsListResponse { func FallbackModelsList() ModelsListResponse {
return ModelsListResponse{Models: DefaultModels()} return ModelsListResponse{Models: DefaultModels()}
} }

View File

@@ -2,7 +2,7 @@ package gemini
import "testing" import "testing"
func TestDefaultModels_ContainsImageModels(t *testing.T) { func TestDefaultModels_ContainsFallbackCatalogModels(t *testing.T) {
t.Parallel() t.Parallel()
models := DefaultModels() models := DefaultModels()
@@ -13,6 +13,7 @@ func TestDefaultModels_ContainsImageModels(t *testing.T) {
required := []string{ required := []string{
"models/gemini-2.5-flash-image", "models/gemini-2.5-flash-image",
"models/gemini-3.1-pro-preview-customtools",
"models/gemini-3.1-flash-image", "models/gemini-3.1-flash-image",
} }
@@ -26,3 +27,17 @@ func TestDefaultModels_ContainsImageModels(t *testing.T) {
} }
} }
} }
func TestHasFallbackModel_RecognizesCustomtoolsModel(t *testing.T) {
t.Parallel()
if !HasFallbackModel("gemini-3.1-pro-preview-customtools") {
t.Fatalf("expected customtools model to exist in fallback catalog")
}
if !HasFallbackModel("models/gemini-3.1-pro-preview-customtools") {
t.Fatalf("expected prefixed customtools model to exist in fallback catalog")
}
if HasFallbackModel("gemini-unknown") {
t.Fatalf("did not expect unknown model to exist in fallback catalog")
}
}