Unverified Commit 095bef95 authored by YanzheL's avatar YanzheL
Browse files

fix(gemini): add customtools fallback metadata

parent 83a16dec
...@@ -2,6 +2,8 @@ ...@@ -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 { ...@@ -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()}
} }
......
...@@ -2,7 +2,7 @@ package gemini ...@@ -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) { ...@@ -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) { ...@@ -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")
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment