mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-10-25 00:56:52 +08:00
fix(gemini): populate token limits in Gemini model catalog responses
- Populate inputTokenLimit and outputTokenLimit in Home Gemini model responses when positive token limits are available, and backfill missing metadata across duplicate Home catalog entries. - Fall back to ContextLength and MaxCompletionTokens when converting registry models for the Gemini handler so catalog entries using generic limit fields expose Gemini token limits. Closes: #6914
This commit is contained in:
@@ -1077,12 +1077,19 @@ func formatHomeGeminiModel(entry homeModelEntry) map[string]any {
|
||||
if displayName == "" {
|
||||
displayName = entry.id
|
||||
}
|
||||
return map[string]any{
|
||||
model := map[string]any{
|
||||
"name": name,
|
||||
"displayName": displayName,
|
||||
"description": displayName,
|
||||
"supportedGenerationMethods": []string{"generateContent"},
|
||||
}
|
||||
if entry.contextLength > 0 {
|
||||
model["inputTokenLimit"] = entry.contextLength
|
||||
}
|
||||
if entry.maxCompletionTokens > 0 {
|
||||
model["outputTokenLimit"] = entry.maxCompletionTokens
|
||||
}
|
||||
return model
|
||||
}
|
||||
|
||||
func homeGeminiModelMatches(entry homeModelEntry, action string) bool {
|
||||
@@ -1176,10 +1183,17 @@ func decodeHomeModels(raw []byte) ([]homeModelEntry, error) {
|
||||
return nil, fmt.Errorf("home models payload contains an error envelope")
|
||||
}
|
||||
|
||||
sections := make([]string, 0, len(bySection))
|
||||
for section := range bySection {
|
||||
sections = append(sections, section)
|
||||
}
|
||||
sort.Strings(sections)
|
||||
|
||||
indexByID := make(map[string]int)
|
||||
out := make([]homeModelEntry, 0, 256)
|
||||
hasModelRecords := false
|
||||
for section, models := range bySection {
|
||||
for _, section := range sections {
|
||||
models := bySection[section]
|
||||
if models == nil {
|
||||
return nil, fmt.Errorf("home models section %q is null", section)
|
||||
}
|
||||
@@ -1201,12 +1215,7 @@ func decodeHomeModels(raw []byte) ([]homeModelEntry, error) {
|
||||
Provider: provider,
|
||||
NativeCapabilities: nativeCapabilities,
|
||||
}
|
||||
if index, ok := indexByID[id]; ok {
|
||||
out[index].providers = appendUniqueHomeProvider(out[index].providers, provider)
|
||||
out[index].nativeCapabilityRoutes = append(out[index].nativeCapabilityRoutes, route)
|
||||
continue
|
||||
}
|
||||
|
||||
created := homeModelInt64Value(model, "created")
|
||||
ownedBy, _ := model["owned_by"].(string)
|
||||
ownedBy = strings.TrimSpace(ownedBy)
|
||||
displayName, _ := model["display_name"].(string)
|
||||
@@ -1215,17 +1224,47 @@ func decodeHomeModels(raw []byte) ([]homeModelEntry, error) {
|
||||
displayName, _ = model["displayName"].(string)
|
||||
displayName = strings.TrimSpace(displayName)
|
||||
}
|
||||
contextLength := int(homeModelInt64Value(model, "context_length", "contextLength", "inputTokenLimit", "max_input_tokens"))
|
||||
maxContextLength := int(homeModelInt64Value(model, "max_context_length", "maxContextLength"))
|
||||
maxCompletionTokens := int(homeModelInt64Value(model, "max_completion_tokens", "maxCompletionTokens", "outputTokenLimit", "max_tokens"))
|
||||
thinking := homeModelThinkingSupport(model)
|
||||
|
||||
if index, ok := indexByID[id]; ok {
|
||||
out[index].providers = appendUniqueHomeProvider(out[index].providers, provider)
|
||||
out[index].nativeCapabilityRoutes = append(out[index].nativeCapabilityRoutes, route)
|
||||
if out[index].created <= 0 && created > 0 {
|
||||
out[index].created = created
|
||||
}
|
||||
if out[index].ownedBy == "" && ownedBy != "" {
|
||||
out[index].ownedBy = ownedBy
|
||||
}
|
||||
if out[index].displayName == "" && displayName != "" {
|
||||
out[index].displayName = displayName
|
||||
}
|
||||
if out[index].contextLength <= 0 && contextLength > 0 {
|
||||
out[index].contextLength = contextLength
|
||||
}
|
||||
if out[index].maxContextLength <= 0 && maxContextLength > 0 {
|
||||
out[index].maxContextLength = maxContextLength
|
||||
}
|
||||
if out[index].maxCompletionTokens <= 0 && maxCompletionTokens > 0 {
|
||||
out[index].maxCompletionTokens = maxCompletionTokens
|
||||
}
|
||||
if out[index].thinking == nil && thinking != nil {
|
||||
out[index].thinking = thinking
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
indexByID[id] = len(out)
|
||||
out = append(out, homeModelEntry{
|
||||
id: id,
|
||||
created: homeModelInt64Value(model, "created"),
|
||||
created: created,
|
||||
ownedBy: ownedBy,
|
||||
displayName: displayName,
|
||||
contextLength: int(homeModelInt64Value(model, "context_length", "contextLength", "inputTokenLimit", "max_input_tokens")),
|
||||
maxContextLength: int(homeModelInt64Value(model, "max_context_length", "maxContextLength")),
|
||||
maxCompletionTokens: int(homeModelInt64Value(model, "max_completion_tokens", "maxCompletionTokens", "outputTokenLimit", "max_tokens")),
|
||||
contextLength: contextLength,
|
||||
maxContextLength: maxContextLength,
|
||||
maxCompletionTokens: maxCompletionTokens,
|
||||
thinking: thinking,
|
||||
providers: appendUniqueHomeProvider(nil, provider),
|
||||
nativeCapabilityRoutes: []registry.NativeCapabilityRoute{route},
|
||||
|
||||
@@ -2814,6 +2814,99 @@ func TestDecodeHomeModelsKeepsTokenMetadata(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeGeminiModelsIncludeTokenLimits(t *testing.T) {
|
||||
withLimits := formatHomeGeminiModel(homeModelEntry{
|
||||
id: "gemini-3.8-flash",
|
||||
displayName: "Gemini 3.8 Flash",
|
||||
contextLength: 1048576,
|
||||
maxCompletionTokens: 65536,
|
||||
})
|
||||
if got := withLimits["name"]; got != "models/gemini-3.8-flash" {
|
||||
t.Fatalf("name = %v, want models/gemini-3.8-flash", got)
|
||||
}
|
||||
if got := withLimits["displayName"]; got != "Gemini 3.8 Flash" {
|
||||
t.Fatalf("displayName = %v, want Gemini 3.8 Flash", got)
|
||||
}
|
||||
if got := withLimits["inputTokenLimit"]; got != 1048576 {
|
||||
t.Fatalf("inputTokenLimit = %v, want 1048576", got)
|
||||
}
|
||||
if got := withLimits["outputTokenLimit"]; got != 65536 {
|
||||
t.Fatalf("outputTokenLimit = %v, want 65536", got)
|
||||
}
|
||||
|
||||
withoutLimits := formatHomeGeminiModel(homeModelEntry{
|
||||
id: "gemini-no-limits",
|
||||
})
|
||||
if _, ok := withoutLimits["inputTokenLimit"]; ok {
|
||||
t.Fatalf("expected inputTokenLimit to be omitted when unknown, got %v", withoutLimits)
|
||||
}
|
||||
if _, ok := withoutLimits["outputTokenLimit"]; ok {
|
||||
t.Fatalf("expected outputTokenLimit to be omitted when unknown, got %v", withoutLimits)
|
||||
}
|
||||
|
||||
server := newTestServer(t)
|
||||
server.cfg.Home.Enabled = true
|
||||
engine := gin.New()
|
||||
engine.GET("/v1beta/models", server.geminiModelsHandler(nil))
|
||||
engine.GET("/v1beta/models/*action", server.geminiGetHandler(nil))
|
||||
|
||||
client := newHomeCatalogClient(t, `{
|
||||
"a-custom": [
|
||||
{
|
||||
"id": "gemini-3.8-flash"
|
||||
}
|
||||
],
|
||||
"gemini": [
|
||||
{
|
||||
"id": "gemini-3.8-flash",
|
||||
"display_name": "Gemini 3.8 Flash",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 65536
|
||||
}
|
||||
]
|
||||
}`)
|
||||
previousHome := home.Current()
|
||||
home.SetCurrent(client)
|
||||
t.Cleanup(func() { home.SetCurrent(previousHome) })
|
||||
|
||||
listRecorder := httptest.NewRecorder()
|
||||
engine.ServeHTTP(listRecorder, httptest.NewRequest(http.MethodGet, "/v1beta/models", nil))
|
||||
if listRecorder.Code != http.StatusOK {
|
||||
t.Fatalf("GET /v1beta/models status = %d; body=%s", listRecorder.Code, listRecorder.Body.String())
|
||||
}
|
||||
var listPayload struct {
|
||||
Models []map[string]any `json:"models"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(listRecorder.Body.Bytes(), &listPayload); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal /v1beta/models response: %v", errUnmarshal)
|
||||
}
|
||||
if len(listPayload.Models) != 1 {
|
||||
t.Fatalf("models len = %d, want 1; body=%s", len(listPayload.Models), listRecorder.Body.String())
|
||||
}
|
||||
if got := homeModelInt64Value(listPayload.Models[0], "inputTokenLimit"); got != 1048576 {
|
||||
t.Fatalf("list inputTokenLimit = %d, want 1048576; model=%v", got, listPayload.Models[0])
|
||||
}
|
||||
if got := homeModelInt64Value(listPayload.Models[0], "outputTokenLimit"); got != 65536 {
|
||||
t.Fatalf("list outputTokenLimit = %d, want 65536; model=%v", got, listPayload.Models[0])
|
||||
}
|
||||
|
||||
getRecorder := httptest.NewRecorder()
|
||||
engine.ServeHTTP(getRecorder, httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-3.8-flash", nil))
|
||||
if getRecorder.Code != http.StatusOK {
|
||||
t.Fatalf("GET /v1beta/models/gemini-3.8-flash status = %d; body=%s", getRecorder.Code, getRecorder.Body.String())
|
||||
}
|
||||
var getPayload map[string]any
|
||||
if errUnmarshal := json.Unmarshal(getRecorder.Body.Bytes(), &getPayload); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal /v1beta/models/gemini-3.8-flash response: %v", errUnmarshal)
|
||||
}
|
||||
if got := homeModelInt64Value(getPayload, "inputTokenLimit"); got != 1048576 {
|
||||
t.Fatalf("get inputTokenLimit = %d, want 1048576; model=%v", got, getPayload)
|
||||
}
|
||||
if got := homeModelInt64Value(getPayload, "outputTokenLimit"); got != 65536 {
|
||||
t.Fatalf("get outputTokenLimit = %d, want 65536; model=%v", got, getPayload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeCodexModels_MaxContextLength(t *testing.T) {
|
||||
entries, errDecode := decodeHomeModels([]byte(`{
|
||||
"codex": [
|
||||
|
||||
@@ -1803,11 +1803,19 @@ func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string)
|
||||
if model.Description != "" {
|
||||
result["description"] = model.Description
|
||||
}
|
||||
if model.InputTokenLimit > 0 {
|
||||
result["inputTokenLimit"] = model.InputTokenLimit
|
||||
inputTokenLimit := model.InputTokenLimit
|
||||
if inputTokenLimit <= 0 {
|
||||
inputTokenLimit = model.ContextLength
|
||||
}
|
||||
if model.OutputTokenLimit > 0 {
|
||||
result["outputTokenLimit"] = model.OutputTokenLimit
|
||||
if inputTokenLimit > 0 {
|
||||
result["inputTokenLimit"] = inputTokenLimit
|
||||
}
|
||||
outputTokenLimit := model.OutputTokenLimit
|
||||
if outputTokenLimit <= 0 {
|
||||
outputTokenLimit = model.MaxCompletionTokens
|
||||
}
|
||||
if outputTokenLimit > 0 {
|
||||
result["outputTokenLimit"] = outputTokenLimit
|
||||
}
|
||||
if len(model.SupportedGenerationMethods) > 0 {
|
||||
result["supportedGenerationMethods"] = append([]string(nil), model.SupportedGenerationMethods...)
|
||||
|
||||
@@ -68,6 +68,70 @@ func TestGetAvailableModelsClaudeIncludesTokenLimits(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAvailableModelsGeminiIncludesTokenLimits(t *testing.T) {
|
||||
r := newTestModelRegistry()
|
||||
r.RegisterClient("client-1", "gemini", []*ModelInfo{
|
||||
{
|
||||
ID: "gemini-3.8-flash",
|
||||
Name: "models/gemini-3.8-flash",
|
||||
DisplayName: "Gemini 3.8 Flash",
|
||||
ContextLength: 1048576,
|
||||
MaxCompletionTokens: 65536,
|
||||
},
|
||||
{
|
||||
ID: "gemini-2.5-pro",
|
||||
Name: "models/gemini-2.5-pro",
|
||||
DisplayName: "Gemini 2.5 Pro",
|
||||
InputTokenLimit: 1048576,
|
||||
OutputTokenLimit: 65536,
|
||||
},
|
||||
{
|
||||
ID: "gemini-no-limits",
|
||||
Name: "models/gemini-no-limits",
|
||||
},
|
||||
})
|
||||
|
||||
models := r.GetAvailableModels("gemini")
|
||||
byName := make(map[string]map[string]any, len(models))
|
||||
for _, m := range models {
|
||||
name, _ := m["name"].(string)
|
||||
byName[name] = m
|
||||
}
|
||||
|
||||
fromContext, ok := byName["models/gemini-3.8-flash"]
|
||||
if !ok {
|
||||
t.Fatalf("expected models/gemini-3.8-flash in available models, got %v", byName)
|
||||
}
|
||||
if got := fromContext["inputTokenLimit"]; got != 1048576 {
|
||||
t.Fatalf("gemini-3.8-flash inputTokenLimit = %v, want 1048576", got)
|
||||
}
|
||||
if got := fromContext["outputTokenLimit"]; got != 65536 {
|
||||
t.Fatalf("gemini-3.8-flash outputTokenLimit = %v, want 65536", got)
|
||||
}
|
||||
|
||||
fromExplicit, ok := byName["models/gemini-2.5-pro"]
|
||||
if !ok {
|
||||
t.Fatalf("expected models/gemini-2.5-pro in available models, got %v", byName)
|
||||
}
|
||||
if got := fromExplicit["inputTokenLimit"]; got != 1048576 {
|
||||
t.Fatalf("gemini-2.5-pro inputTokenLimit = %v, want 1048576", got)
|
||||
}
|
||||
if got := fromExplicit["outputTokenLimit"]; got != 65536 {
|
||||
t.Fatalf("gemini-2.5-pro outputTokenLimit = %v, want 65536", got)
|
||||
}
|
||||
|
||||
noLimits, ok := byName["models/gemini-no-limits"]
|
||||
if !ok {
|
||||
t.Fatalf("expected models/gemini-no-limits in available models, got %v", byName)
|
||||
}
|
||||
if _, exists := noLimits["inputTokenLimit"]; exists {
|
||||
t.Fatalf("expected inputTokenLimit to be omitted when unknown, got %v", noLimits)
|
||||
}
|
||||
if _, exists := noLimits["outputTokenLimit"]; exists {
|
||||
t.Fatalf("expected outputTokenLimit to be omitted when unknown, got %v", noLimits)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAvailableModelsInvalidatesCacheOnRegistryChanges(t *testing.T) {
|
||||
r := newTestModelRegistry()
|
||||
r.RegisterClient("client-1", "OpenAI", []*ModelInfo{{ID: "m1", OwnedBy: "team-a", DisplayName: "Model One"}})
|
||||
|
||||
Reference in New Issue
Block a user