diff --git a/config.example.yaml b/config.example.yaml index 73e2a8ac..2a35fe68 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -90,6 +90,9 @@ ws-auth: false # headers: # X-Custom-Header: "custom-value" # proxy-url: "socks5://proxy.example.com:1080" +# models: +# - name: "gemini-2.5-flash" # upstream model name +# alias: "gemini-flash" # client alias mapped to the upstream model # excluded-models: # - "gemini-2.5-pro" # exclude specific models from this provider (exact match) # - "gemini-2.5-*" # wildcard matching prefix (e.g. gemini-2.5-flash, gemini-2.5-pro) @@ -106,7 +109,7 @@ ws-auth: false # X-Custom-Header: "custom-value" # proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override # models: -# - name: "gpt-5-codex" # upstream model name +# - name: "gpt-5-codex" # upstream model name # alias: "codex-latest" # client alias mapped to the upstream model # excluded-models: # - "gpt-5.1" # exclude specific models (exact match) @@ -125,7 +128,7 @@ ws-auth: false # proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override # models: # - name: "claude-3-5-sonnet-20241022" # upstream model name -# alias: "claude-sonnet-latest" # client alias mapped to the upstream model +# alias: "claude-sonnet-latest" # client alias mapped to the upstream model # excluded-models: # - "claude-opus-4-5-20251101" # exclude specific models (exact match) # - "claude-3-*" # wildcard matching prefix (e.g. claude-3-7-sonnet-20250219) diff --git a/internal/config/config.go b/internal/config/config.go index 760be600..0cde69c7 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -318,6 +318,9 @@ type GeminiKey struct { // ProxyURL optionally overrides the global proxy for this API key. ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"` + // Models defines upstream model names and aliases for request routing. + Models []GeminiModel `yaml:"models,omitempty" json:"models,omitempty"` + // Headers optionally adds extra HTTP headers for requests sent with this key. Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` @@ -325,6 +328,15 @@ type GeminiKey struct { ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` } +// GeminiModel describes a mapping between an alias and the actual upstream model name. +type GeminiModel struct { + // Name is the upstream model identifier used when issuing requests. + Name string `yaml:"name" json:"name"` + + // Alias is the client-facing model name that maps to Name. + Alias string `yaml:"alias" json:"alias"` +} + // OpenAICompatibility represents the configuration for OpenAI API compatibility // with external providers, allowing model aliases to be routed through OpenAI API format. type OpenAICompatibility struct { diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index f211ba62..da57150d 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -78,6 +78,13 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r defer reporter.trackFailure(ctx, &err) upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata) + if modelOverride := e.resolveUpstreamModel(upstreamModel, auth); modelOverride != "" { + upstreamModel = modelOverride + } else if !strings.EqualFold(upstreamModel, req.Model) { + if modelOverride := e.resolveUpstreamModel(req.Model, auth); modelOverride != "" { + upstreamModel = modelOverride + } + } // Official Gemini API via API key or OAuth bearer from := opts.SourceFormat @@ -174,6 +181,13 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A defer reporter.trackFailure(ctx, &err) upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata) + if modelOverride := e.resolveUpstreamModel(upstreamModel, auth); modelOverride != "" { + upstreamModel = modelOverride + } else if !strings.EqualFold(upstreamModel, req.Model) { + if modelOverride := e.resolveUpstreamModel(req.Model, auth); modelOverride != "" { + upstreamModel = modelOverride + } + } from := opts.SourceFormat to := sdktranslator.FromString("gemini") @@ -287,6 +301,15 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { apiKey, bearer := geminiCreds(auth) + upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata) + if modelOverride := e.resolveUpstreamModel(upstreamModel, auth); modelOverride != "" { + upstreamModel = modelOverride + } else if !strings.EqualFold(upstreamModel, req.Model) { + if modelOverride := e.resolveUpstreamModel(req.Model, auth); modelOverride != "" { + upstreamModel = modelOverride + } + } + from := opts.SourceFormat to := sdktranslator.FromString("gemini") translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false) @@ -297,9 +320,10 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools") translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig") translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings") + translatedReq, _ = sjson.SetBytes(translatedReq, "model", upstreamModel) baseURL := resolveGeminiBaseURL(auth) - url := fmt.Sprintf("%s/%s/models/%s:%s", baseURL, glAPIVersion, req.Model, "countTokens") + url := fmt.Sprintf("%s/%s/models/%s:%s", baseURL, glAPIVersion, upstreamModel, "countTokens") requestBody := bytes.NewReader(translatedReq) @@ -398,6 +422,90 @@ func resolveGeminiBaseURL(auth *cliproxyauth.Auth) string { return base } +func (e *GeminiExecutor) resolveUpstreamModel(alias string, auth *cliproxyauth.Auth) string { + trimmed := strings.TrimSpace(alias) + if trimmed == "" { + return "" + } + + entry := e.resolveGeminiConfig(auth) + if entry == nil { + return "" + } + + normalizedModel, metadata := util.NormalizeThinkingModel(trimmed) + + // Candidate names to match against configured aliases/names. + candidates := []string{strings.TrimSpace(normalizedModel)} + if !strings.EqualFold(normalizedModel, trimmed) { + candidates = append(candidates, trimmed) + } + if original := util.ResolveOriginalModel(normalizedModel, metadata); original != "" && !strings.EqualFold(original, normalizedModel) { + candidates = append(candidates, original) + } + + for i := range entry.Models { + model := entry.Models[i] + name := strings.TrimSpace(model.Name) + modelAlias := strings.TrimSpace(model.Alias) + + for _, candidate := range candidates { + if candidate == "" { + continue + } + if modelAlias != "" && strings.EqualFold(modelAlias, candidate) { + if name != "" { + return name + } + return candidate + } + if name != "" && strings.EqualFold(name, candidate) { + return name + } + } + } + return "" +} + +func (e *GeminiExecutor) resolveGeminiConfig(auth *cliproxyauth.Auth) *config.GeminiKey { + if auth == nil || e.cfg == nil { + return nil + } + var attrKey, attrBase string + if auth.Attributes != nil { + attrKey = strings.TrimSpace(auth.Attributes["api_key"]) + attrBase = strings.TrimSpace(auth.Attributes["base_url"]) + } + for i := range e.cfg.GeminiKey { + entry := &e.cfg.GeminiKey[i] + cfgKey := strings.TrimSpace(entry.APIKey) + cfgBase := strings.TrimSpace(entry.BaseURL) + if attrKey != "" && attrBase != "" { + if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) { + return entry + } + continue + } + if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { + if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey != "" { + for i := range e.cfg.GeminiKey { + entry := &e.cfg.GeminiKey[i] + if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) { + return entry + } + } + } + return nil +} + func applyGeminiHeaders(req *http.Request, auth *cliproxyauth.Auth) { var attrs map[string]string if auth != nil { diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index ae56e4b6..4101bf22 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -710,6 +710,9 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) { case "gemini": models = registry.GetGeminiModels() if entry := s.resolveConfigGeminiKey(a); entry != nil { + if len(entry.Models) > 0 { + models = buildGeminiConfigModels(entry) + } if authKind == "apikey" { excluded = entry.ExcludedModels } @@ -1146,7 +1149,7 @@ func buildVertexCompatConfigModels(entry *config.VertexCompatKey) []*ModelInfo { ID: alias, Object: "model", Created: now, - OwnedBy: "vertex", + OwnedBy: "google", Type: "vertex", DisplayName: display, }) @@ -1241,6 +1244,44 @@ func applyOAuthModelMappings(cfg *config.Config, provider, authKind string, mode return out } +func buildGeminiConfigModels(entry *config.GeminiKey) []*ModelInfo { + if entry == nil || len(entry.Models) == 0 { + return nil + } + now := time.Now().Unix() + out := make([]*ModelInfo, 0, len(entry.Models)) + seen := make(map[string]struct{}, len(entry.Models)) + for i := range entry.Models { + model := entry.Models[i] + name := strings.TrimSpace(model.Name) + alias := strings.TrimSpace(model.Alias) + if alias == "" { + alias = name + } + if alias == "" { + continue + } + key := strings.ToLower(alias) + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + display := name + if display == "" { + display = alias + } + out = append(out, &ModelInfo{ + ID: alias, + Object: "model", + Created: now, + OwnedBy: "google", + Type: "gemini", + DisplayName: display, + }) + } + return out +} + func buildClaudeConfigModels(entry *config.ClaudeKey) []*ModelInfo { if entry == nil || len(entry.Models) == 0 { return nil @@ -1271,7 +1312,7 @@ func buildClaudeConfigModels(entry *config.ClaudeKey) []*ModelInfo { ID: alias, Object: "model", Created: now, - OwnedBy: "claude", + OwnedBy: "anthropic", Type: "claude", DisplayName: display, })