Skip to content

Commit d6272d3

Browse files
luispaterluispater
authored
Merge pull request router-for-me#177 from router-for-me/aistudio
feat(registry): unify Gemini models and add AI Studio set
2 parents 5891785 + c99d0df commit d6272d3

File tree

5 files changed

+81
-142
lines changed

5 files changed

+81
-142
lines changed

‎internal/api/server.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
225225
envManagementSecret := envAdminPasswordSet && envAdminPassword != ""
226226

227227
// Create server instance
228-
providerNames := make([]string, 0, len(cfg.OpenAICompatibility))
228+
providerNames := make([]string, 0, len(cfg.OpenAICompatibility))
229229
for _, p := range cfg.OpenAICompatibility {
230230
providerNames = append(providerNames, p.Name)
231231
}
@@ -914,5 +914,3 @@ func AuthMiddleware(manager *sdkaccess.Manager) gin.HandlerFunc {
914914
}
915915
}
916916
}
917-
918-
// legacy clientsToSlice removed; handlers no longer consume legacy client slices

‎internal/registry/model_definitions.go‎

Lines changed: 31 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,8 @@ func GetClaudeModels() []*ModelInfo {
6868
}
6969
}
7070

71-
// GetGeminiModels returns the standard Gemini model definitions
72-
func GetGeminiModels() []*ModelInfo {
71+
// GeminiModels returns the shared base Gemini model set used by multiple providers.
72+
func GeminiModels() []*ModelInfo {
7373
return []*ModelInfo{
7474
{
7575
ID: "gemini-2.5-flash",
@@ -108,7 +108,7 @@ func GetGeminiModels() []*ModelInfo {
108108
Name: "models/gemini-2.5-flash-lite",
109109
Version: "2.5",
110110
DisplayName: "Gemini 2.5 Flash Lite",
111-
Description: "Stable release (June 17th, 2025) of Gemini 2.5 Flash Lite",
111+
Description: "Our smallest and most cost effective model, built for at scale usage.",
112112
InputTokenLimit: 1048576,
113113
OutputTokenLimit: 65536,
114114
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
@@ -144,80 +144,61 @@ func GetGeminiModels() []*ModelInfo {
144144
}
145145
}
146146

147+
// GetGeminiModels returns the standard Gemini model definitions
148+
func GetGeminiModels() []*ModelInfo { return GeminiModels() }
149+
147150
// GetGeminiCLIModels returns the standard Gemini model definitions
148-
func GetGeminiCLIModels() []*ModelInfo {
149-
return []*ModelInfo{
150-
{
151-
ID: "gemini-2.5-flash",
152-
Object: "model",
153-
Created: time.Now().Unix(),
154-
OwnedBy: "google",
155-
Type: "gemini",
156-
Name: "models/gemini-2.5-flash",
157-
Version: "001",
158-
DisplayName: "Gemini 2.5 Flash",
159-
Description: "Stable version of Gemini 2.5 Flash, our mid-size multimodal model that supports up to 1 million tokens, released in June of 2025.",
160-
InputTokenLimit: 1048576,
161-
OutputTokenLimit: 65536,
162-
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
163-
},
164-
{
165-
ID: "gemini-2.5-pro",
151+
func GetGeminiCLIModels() []*ModelInfo { return GeminiModels() }
152+
153+
// GetAIStudioModels returns the Gemini model definitions for AI Studio integrations
154+
func GetAIStudioModels() []*ModelInfo {
155+
models := make([]*ModelInfo, 0, 8)
156+
models = append(models, GeminiModels()...)
157+
models = append(models,
158+
&ModelInfo{
159+
ID: "gemini-pro-latest",
166160
Object: "model",
167161
Created: time.Now().Unix(),
168162
OwnedBy: "google",
169163
Type: "gemini",
170-
Name: "models/gemini-2.5-pro",
164+
Name: "models/gemini-pro-latest",
171165
Version: "2.5",
172-
DisplayName: "Gemini 2.5 Pro",
173-
Description: "Stable release (June 17th, 2025) of Gemini 2.5 Pro",
166+
DisplayName: "Gemini Pro Latest",
167+
Description: "Latest release of Gemini Pro",
174168
InputTokenLimit: 1048576,
175169
OutputTokenLimit: 65536,
176170
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
177171
},
178-
{
179-
ID: "gemini-2.5-flash-lite",
172+
&ModelInfo{
173+
ID: "gemini-flash-latest",
180174
Object: "model",
181175
Created: time.Now().Unix(),
182176
OwnedBy: "google",
183177
Type: "gemini",
184-
Name: "models/gemini-2.5-flash-lite",
178+
Name: "models/gemini-flash-latest",
185179
Version: "2.5",
186-
DisplayName: "Gemini 2.5 Flash Lite",
187-
Description: "Our smallest and most cost effective model, built for at scale usage.",
180+
DisplayName: "Gemini Flash Latest",
181+
Description: "Latest release of Gemini Flash",
188182
InputTokenLimit: 1048576,
189183
OutputTokenLimit: 65536,
190184
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
191185
},
192-
{
193-
ID: "gemini-2.5-flash-image-preview",
194-
Object: "model",
195-
Created: time.Now().Unix(),
196-
OwnedBy: "google",
197-
Type: "gemini",
198-
Name: "models/gemini-2.5-flash-image-preview",
199-
Version: "2.5",
200-
DisplayName: "Gemini 2.5 Flash Image Preview",
201-
Description: "State-of-the-art image generation and editing model.",
202-
InputTokenLimit: 1048576,
203-
OutputTokenLimit: 8192,
204-
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
205-
},
206-
{
207-
ID: "gemini-2.5-flash-image",
186+
&ModelInfo{
187+
ID: "gemini-flash-lite-latest",
208188
Object: "model",
209189
Created: time.Now().Unix(),
210190
OwnedBy: "google",
211191
Type: "gemini",
212-
Name: "models/gemini-2.5-flash-image",
192+
Name: "models/gemini-flash-lite-latest",
213193
Version: "2.5",
214-
DisplayName: "Gemini 2.5 Flash Image",
215-
Description: "State-of-the-art image generation and editing model.",
194+
DisplayName: "Gemini Flash-Lite Latest",
195+
Description: "Latest release of Gemini Flash-Lite",
216196
InputTokenLimit: 1048576,
217-
OutputTokenLimit: 8192,
197+
OutputTokenLimit: 65536,
218198
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
219199
},
220-
}
200+
)
201+
return models
221202
}
222203

223204
// GetOpenAIModels returns the standard OpenAI model definitions
@@ -417,7 +398,6 @@ func GetIFlowModels() []*ModelInfo {
417398
{ID: "qwen3-vl-plus", DisplayName: "Qwen3-VL-Plus", Description: "Qwen3 multimodal vision-language"},
418399
{ID: "qwen3-max-preview", DisplayName: "Qwen3-Max-Preview", Description: "Qwen3 Max preview build"},
419400
{ID: "kimi-k2-0905", DisplayName: "Kimi-K2-Instruct-0905", Description: "Moonshot Kimi K2 instruct 0905"},
420-
{ID: "glm-4.5", DisplayName: "GLM-4.5", Description: "Zhipu GLM 4.5 general model"},
421401
{ID: "glm-4.6", DisplayName: "GLM-4.6", Description: "Zhipu GLM 4.6 general model"},
422402
{ID: "kimi-k2", DisplayName: "Kimi-K2", Description: "Moonshot Kimi K2 general model"},
423403
{ID: "deepseek-v3.2", DisplayName: "DeepSeek-V3.2-Exp", Description: "DeepSeek V3.2 experimental"},

‎internal/runtime/executor/aistudio_executor.go‎

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -19,27 +19,27 @@ import (
1919
"github.com/tidwall/sjson"
2020
)
2121

22-
// AistudioExecutor routes AI Studio requests through a websocket-backed transport.
23-
type AistudioExecutor struct {
22+
// AIStudioExecutor routes AI Studio requests through a websocket-backed transport.
23+
type AIStudioExecutor struct {
2424
provider string
2525
relay *wsrelay.Manager
2626
cfg *config.Config
2727
}
2828

29-
// NewAistudioExecutor constructs a websocket executor for the provider name.
30-
func NewAistudioExecutor(cfg *config.Config, provider string, relay *wsrelay.Manager) *AistudioExecutor {
31-
return &AistudioExecutor{provider: strings.ToLower(provider), relay: relay, cfg: cfg}
29+
// NewAIStudioExecutor constructs a websocket executor for the provider name.
30+
func NewAIStudioExecutor(cfg *config.Config, provider string, relay *wsrelay.Manager) *AIStudioExecutor {
31+
return &AIStudioExecutor{provider: strings.ToLower(provider), relay: relay, cfg: cfg}
3232
}
3333

34-
// Identifier returns the provider key served by this executor.
35-
func (e *AistudioExecutor) Identifier() string { return e.provider }
34+
// Identifier returns the logical provider key for routing.
35+
func (e *AIStudioExecutor) Identifier() string { return "aistudio" }
3636

3737
// PrepareRequest is a no-op because websocket transport already injects headers.
38-
func (e *AistudioExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error {
38+
func (e *AIStudioExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error {
3939
return nil
4040
}
4141

42-
func (e *AistudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
42+
func (e *AIStudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
4343
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
4444
defer reporter.trackFailure(ctx, &err)
4545

@@ -66,14 +66,14 @@ func (e *AistudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth,
6666
Method: http.MethodPost,
6767
Headers: wsReq.Headers.Clone(),
6868
Body: bytes.Clone(body.payload),
69-
Provider: e.provider,
69+
Provider: e.Identifier(),
7070
AuthID: authID,
7171
AuthLabel: authLabel,
7272
AuthType: authType,
7373
AuthValue: authValue,
7474
})
7575

76-
wsResp, err := e.relay.NonStream(ctx, e.provider, wsReq)
76+
wsResp, err := e.relay.NonStream(ctx, authID, wsReq)
7777
if err != nil {
7878
recordAPIResponseError(ctx, e.cfg, err)
7979
return resp, err
@@ -92,7 +92,7 @@ func (e *AistudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth,
9292
return resp, nil
9393
}
9494

95-
func (e *AistudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
95+
func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
9696
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
9797
defer reporter.trackFailure(ctx, &err)
9898

@@ -118,13 +118,13 @@ func (e *AistudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
118118
Method: http.MethodPost,
119119
Headers: wsReq.Headers.Clone(),
120120
Body: bytes.Clone(body.payload),
121-
Provider: e.provider,
121+
Provider: e.Identifier(),
122122
AuthID: authID,
123123
AuthLabel: authLabel,
124124
AuthType: authType,
125125
AuthValue: authValue,
126126
})
127-
wsStream, err := e.relay.Stream(ctx, e.provider, wsReq)
127+
wsStream, err := e.relay.Stream(ctx, authID, wsReq)
128128
if err != nil {
129129
recordAPIResponseError(ctx, e.cfg, err)
130130
return nil, err
@@ -151,7 +151,7 @@ func (e *AistudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
151151
case wsrelay.MessageTypeStreamChunk:
152152
if len(event.Payload) > 0 {
153153
appendAPIResponseChunk(ctx, e.cfg, bytes.Clone(event.Payload))
154-
filtered := filterAistudioUsageMetadata(event.Payload)
154+
filtered := filterAIStudioUsageMetadata(event.Payload)
155155
if detail, ok := parseGeminiStreamUsage(filtered); ok {
156156
reporter.publish(ctx, detail)
157157
}
@@ -188,7 +188,7 @@ func (e *AistudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
188188
return stream, nil
189189
}
190190

191-
func (e *AistudioExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
191+
func (e *AIStudioExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
192192
_, body, err := e.translateRequest(req, opts, false)
193193
if err != nil {
194194
return cliproxyexecutor.Response{}, err
@@ -215,13 +215,13 @@ func (e *AistudioExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.A
215215
Method: http.MethodPost,
216216
Headers: wsReq.Headers.Clone(),
217217
Body: bytes.Clone(body.payload),
218-
Provider: e.provider,
218+
Provider: e.Identifier(),
219219
AuthID: authID,
220220
AuthLabel: authLabel,
221221
AuthType: authType,
222222
AuthValue: authValue,
223223
})
224-
resp, err := e.relay.NonStream(ctx, e.provider, wsReq)
224+
resp, err := e.relay.NonStream(ctx, authID, wsReq)
225225
if err != nil {
226226
recordAPIResponseError(ctx, e.cfg, err)
227227
return cliproxyexecutor.Response{}, err
@@ -241,7 +241,7 @@ func (e *AistudioExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.A
241241
return cliproxyexecutor.Response{Payload: []byte(translated)}, nil
242242
}
243243

244-
func (e *AistudioExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
244+
func (e *AIStudioExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
245245
_ = ctx
246246
return auth, nil
247247
}
@@ -252,7 +252,7 @@ type translatedPayload struct {
252252
toFormat sdktranslator.Format
253253
}
254254

255-
func (e *AistudioExecutor) translateRequest(req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) ([]byte, translatedPayload, error) {
255+
func (e *AIStudioExecutor) translateRequest(req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) ([]byte, translatedPayload, error) {
256256
from := opts.SourceFormat
257257
to := sdktranslator.FromString("gemini")
258258
payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), stream)
@@ -275,7 +275,7 @@ func (e *AistudioExecutor) translateRequest(req cliproxyexecutor.Request, opts c
275275
return payload, translatedPayload{payload: payload, action: action, toFormat: to}, nil
276276
}
277277

278-
func (e *AistudioExecutor) buildEndpoint(model, action, alt string) string {
278+
func (e *AIStudioExecutor) buildEndpoint(model, action, alt string) string {
279279
base := fmt.Sprintf("%s/%s/models/%s:%s", glEndpoint, glAPIVersion, model, action)
280280
if action == "streamGenerateContent" {
281281
if alt == "" {
@@ -289,9 +289,9 @@ func (e *AistudioExecutor) buildEndpoint(model, action, alt string) string {
289289
return base
290290
}
291291

292-
// filterAistudioUsageMetadata removes usageMetadata from intermediate SSE events so that
292+
// filterAIStudioUsageMetadata removes usageMetadata from intermediate SSE events so that
293293
// only the terminal chunk retains token statistics.
294-
func filterAistudioUsageMetadata(payload []byte) []byte {
294+
func filterAIStudioUsageMetadata(payload []byte) []byte {
295295
if len(payload) == 0 {
296296
return payload
297297
}

‎sdk/cliproxy/auth/types.go‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -157,15 +157,6 @@ func (a *Auth) AccountInfo() (string, string) {
157157
return "oauth", v
158158
}
159159
}
160-
if strings.HasPrefix(strings.ToLower(strings.TrimSpace(a.Provider)), "aistudio-") {
161-
if label := strings.TrimSpace(a.Label); label != "" {
162-
return "oauth", label
163-
}
164-
if id := strings.TrimSpace(a.ID); id != "" {
165-
return "oauth", id
166-
}
167-
return "oauth", "aistudio"
168-
}
169160
if a.Attributes != nil {
170161
if v := a.Attributes["api_key"]; v != "" {
171162
return "api_key", v

0 commit comments

Comments
 (0)