refactor(thinking): extract antigravity logic into a dedicated provider

This commit is contained in:
hkfires
2026-01-15 19:08:22 +08:00
parent fe5b3c80cb
commit 4ad6189487
6 changed files with 205 additions and 90 deletions

View File

@@ -55,23 +55,15 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
} }
levelStr := fmt.Sprintf("%-5s", level) levelStr := fmt.Sprintf("%-5s", level)
// Build fields string (excluding request_id which is already shown) // Build fields string (only print fields in logFieldOrder)
var fieldsStr string var fieldsStr string
if len(entry.Data) > 0 { if len(entry.Data) > 0 {
seen := make(map[string]bool)
var fields []string var fields []string
for _, k := range logFieldOrder { for _, k := range logFieldOrder {
if v, ok := entry.Data[k]; ok { if v, ok := entry.Data[k]; ok {
fields = append(fields, fmt.Sprintf("%s=%v", k, v)) fields = append(fields, fmt.Sprintf("%s=%v", k, v))
seen[k] = true
} }
} }
for k, v := range entry.Data {
if k == "request_id" || seen[k] {
continue
}
fields = append(fields, fmt.Sprintf("%s=%v", k, v))
}
if len(fields) > 0 { if len(fields) > 0 {
fieldsStr = " " + strings.Join(fields, " ") fieldsStr = " " + strings.Join(fields, " ")
} }

View File

@@ -141,8 +141,6 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
return resp, err return resp, err
} }
// Preserve Claude special handling (use baseModel for registry lookups)
translated = normalizeAntigravityThinking(baseModel, translated, isClaude)
translated = applyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated) translated = applyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated)
baseURLs := antigravityBaseURLFallbackOrder(auth) baseURLs := antigravityBaseURLFallbackOrder(auth)
@@ -262,8 +260,6 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
return resp, err return resp, err
} }
// Preserve Claude special handling (use baseModel for registry lookups)
translated = normalizeAntigravityThinking(baseModel, translated, true)
translated = applyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated) translated = applyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated)
baseURLs := antigravityBaseURLFallbackOrder(auth) baseURLs := antigravityBaseURLFallbackOrder(auth)
@@ -603,7 +599,6 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
baseModel := thinking.ParseSuffix(req.Model).ModelName baseModel := thinking.ParseSuffix(req.Model).ModelName
ctx = context.WithValue(ctx, "alt", "") ctx = context.WithValue(ctx, "alt", "")
isClaude := strings.Contains(strings.ToLower(baseModel), "claude")
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
if errToken != nil { if errToken != nil {
@@ -631,8 +626,6 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
return nil, err return nil, err
} }
// Preserve Claude special handling (use baseModel for registry lookups)
translated = normalizeAntigravityThinking(baseModel, translated, isClaude)
translated = applyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated) translated = applyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated)
baseURLs := antigravityBaseURLFallbackOrder(auth) baseURLs := antigravityBaseURLFallbackOrder(auth)
@@ -790,7 +783,6 @@ func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Au
// CountTokens counts tokens for the given request using the Antigravity API. // CountTokens counts tokens for the given request using the Antigravity API.
func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
baseModel := thinking.ParseSuffix(req.Model).ModelName baseModel := thinking.ParseSuffix(req.Model).ModelName
isClaude := strings.Contains(strings.ToLower(baseModel), "claude")
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
if errToken != nil { if errToken != nil {
@@ -815,8 +807,6 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut
return cliproxyexecutor.Response{}, err return cliproxyexecutor.Response{}, err
} }
// Preserve Claude special handling (use baseModel for registry lookups)
payload = normalizeAntigravityThinking(baseModel, payload, isClaude)
payload = deleteJSONField(payload, "project") payload = deleteJSONField(payload, "project")
payload = deleteJSONField(payload, "model") payload = deleteJSONField(payload, "model")
payload = deleteJSONField(payload, "request.safetySettings") payload = deleteJSONField(payload, "request.safetySettings")
@@ -1447,71 +1437,3 @@ func generateProjectID() string {
randomPart := strings.ToLower(uuid.NewString())[:5] randomPart := strings.ToLower(uuid.NewString())[:5]
return adj + "-" + noun + "-" + randomPart return adj + "-" + noun + "-" + randomPart
} }
// normalizeAntigravityThinking performs Antigravity-specific thinking config normalization.
// This function is called AFTER thinking.ApplyThinking() to apply Claude-specific constraints.
//
// It handles:
// - Stripping thinking config for unsupported models
// - Normalizing budget to model range (via thinking.ClampBudget)
// - For Claude models: ensuring thinking budget < max_tokens
// - For Claude models: removing thinkingConfig if budget < minimum allowed
func normalizeAntigravityThinking(model string, payload []byte, isClaude bool) []byte {
modelInfo := registry.LookupModelInfo(model)
if modelInfo == nil || modelInfo.Thinking == nil {
// Model doesn't support thinking - strip any thinking config
return thinking.StripThinkingConfig(payload, "antigravity")
}
budget := gjson.GetBytes(payload, "request.generationConfig.thinkingConfig.thinkingBudget")
if !budget.Exists() {
return payload
}
raw := int(budget.Int())
normalized := thinking.ClampBudget(raw, modelInfo, "antigravity")
if isClaude {
effectiveMax, setDefaultMax := antigravityEffectiveMaxTokens(model, payload)
if effectiveMax > 0 && normalized >= effectiveMax {
normalized = effectiveMax - 1
}
minBudget := antigravityMinThinkingBudget(model)
if minBudget > 0 && normalized >= 0 && normalized < minBudget {
// Budget is below minimum, remove thinking config entirely
payload, _ = sjson.DeleteBytes(payload, "request.generationConfig.thinkingConfig")
return payload
}
if setDefaultMax {
if res, errSet := sjson.SetBytes(payload, "request.generationConfig.maxOutputTokens", effectiveMax); errSet == nil {
payload = res
}
}
}
updated, err := sjson.SetBytes(payload, "request.generationConfig.thinkingConfig.thinkingBudget", normalized)
if err != nil {
return payload
}
return updated
}
// antigravityEffectiveMaxTokens returns the max tokens to cap thinking:
// prefer request-provided maxOutputTokens; otherwise fall back to model default.
// The boolean indicates whether the value came from the model default (and thus should be written back).
func antigravityEffectiveMaxTokens(model string, payload []byte) (max int, fromModel bool) {
if maxTok := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxTok.Exists() && maxTok.Int() > 0 {
return int(maxTok.Int()), false
}
if modelInfo := registry.LookupModelInfo(model); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 {
return modelInfo.MaxCompletionTokens, true
}
return 0, false
}
// antigravityMinThinkingBudget returns the minimum thinking budget for a model.
// Falls back to -1 if no model info is found.
func antigravityMinThinkingBudget(model string) int {
if modelInfo := registry.LookupModelInfo(model); modelInfo != nil && modelInfo.Thinking != nil {
return modelInfo.Thinking.Min
}
return -1
}

View File

@@ -1,6 +1,7 @@
package executor package executor
import ( import (
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/antigravity"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/claude" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/claude"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/codex" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/codex"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/gemini" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/gemini"

View File

@@ -0,0 +1,201 @@
// Package antigravity implements thinking configuration for Antigravity API format.
//
// Antigravity uses request.generationConfig.thinkingConfig.* path (same as gemini-cli)
// but requires additional normalization for Claude models:
// - Ensure thinking budget < max_tokens
// - Remove thinkingConfig if budget < minimum allowed
package antigravity
import (
"strings"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// Applier applies thinking configuration for Antigravity API format.
type Applier struct{}
var _ thinking.ProviderApplier = (*Applier)(nil)
// NewApplier creates a new Antigravity thinking applier.
func NewApplier() *Applier {
return &Applier{}
}
func init() {
thinking.RegisterProvider("antigravity", NewApplier())
}
// Apply applies thinking configuration to Antigravity request body.
//
// For Claude models, additional constraints are applied:
// - Ensure thinking budget < max_tokens
// - Remove thinkingConfig if budget < minimum allowed
func (a *Applier) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) {
if thinking.IsUserDefinedModel(modelInfo) {
return a.applyCompatible(body, config, modelInfo)
}
if modelInfo.Thinking == nil {
return body, nil
}
if config.Mode != thinking.ModeBudget && config.Mode != thinking.ModeLevel && config.Mode != thinking.ModeNone && config.Mode != thinking.ModeAuto {
return body, nil
}
if len(body) == 0 || !gjson.ValidBytes(body) {
body = []byte(`{}`)
}
isClaude := strings.Contains(strings.ToLower(modelInfo.ID), "claude")
// ModeAuto: Always use Budget format with thinkingBudget=-1
if config.Mode == thinking.ModeAuto {
return a.applyBudgetFormat(body, config, modelInfo, isClaude)
}
if config.Mode == thinking.ModeBudget {
return a.applyBudgetFormat(body, config, modelInfo, isClaude)
}
// For non-auto modes, choose format based on model capabilities
support := modelInfo.Thinking
if len(support.Levels) > 0 {
return a.applyLevelFormat(body, config)
}
return a.applyBudgetFormat(body, config, modelInfo, isClaude)
}
func (a *Applier) applyCompatible(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) {
if config.Mode != thinking.ModeBudget && config.Mode != thinking.ModeLevel && config.Mode != thinking.ModeNone && config.Mode != thinking.ModeAuto {
return body, nil
}
if len(body) == 0 || !gjson.ValidBytes(body) {
body = []byte(`{}`)
}
isClaude := false
if modelInfo != nil {
isClaude = strings.Contains(strings.ToLower(modelInfo.ID), "claude")
}
if config.Mode == thinking.ModeAuto {
return a.applyBudgetFormat(body, config, modelInfo, isClaude)
}
if config.Mode == thinking.ModeLevel || (config.Mode == thinking.ModeNone && config.Level != "") {
return a.applyLevelFormat(body, config)
}
return a.applyBudgetFormat(body, config, modelInfo, isClaude)
}
func (a *Applier) applyLevelFormat(body []byte, config thinking.ThinkingConfig) ([]byte, error) {
// Remove conflicting field to avoid both thinkingLevel and thinkingBudget in output
result, _ := sjson.DeleteBytes(body, "request.generationConfig.thinkingConfig.thinkingBudget")
// Normalize includeThoughts field name to avoid oneof conflicts in upstream JSON parsing.
result, _ = sjson.DeleteBytes(result, "request.generationConfig.thinkingConfig.include_thoughts")
if config.Mode == thinking.ModeNone {
result, _ = sjson.SetBytes(result, "request.generationConfig.thinkingConfig.includeThoughts", false)
if config.Level != "" {
result, _ = sjson.SetBytes(result, "request.generationConfig.thinkingConfig.thinkingLevel", string(config.Level))
}
return result, nil
}
// Only handle ModeLevel - budget conversion should be done by upper layer
if config.Mode != thinking.ModeLevel {
return body, nil
}
level := string(config.Level)
result, _ = sjson.SetBytes(result, "request.generationConfig.thinkingConfig.thinkingLevel", level)
result, _ = sjson.SetBytes(result, "request.generationConfig.thinkingConfig.includeThoughts", true)
return result, nil
}
func (a *Applier) applyBudgetFormat(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo, isClaude bool) ([]byte, error) {
// Remove conflicting field to avoid both thinkingLevel and thinkingBudget in output
result, _ := sjson.DeleteBytes(body, "request.generationConfig.thinkingConfig.thinkingLevel")
// Normalize includeThoughts field name to avoid oneof conflicts in upstream JSON parsing.
result, _ = sjson.DeleteBytes(result, "request.generationConfig.thinkingConfig.include_thoughts")
budget := config.Budget
includeThoughts := false
switch config.Mode {
case thinking.ModeNone:
includeThoughts = false
case thinking.ModeAuto:
includeThoughts = true
default:
includeThoughts = budget > 0
}
// Apply Claude-specific constraints
if isClaude && modelInfo != nil {
budget, result = a.normalizeClaudeBudget(budget, result, modelInfo)
// Check if budget was removed entirely
if budget == -2 {
return result, nil
}
}
result, _ = sjson.SetBytes(result, "request.generationConfig.thinkingConfig.thinkingBudget", budget)
result, _ = sjson.SetBytes(result, "request.generationConfig.thinkingConfig.includeThoughts", includeThoughts)
return result, nil
}
// normalizeClaudeBudget applies Claude-specific constraints to thinking budget.
//
// It handles:
// - Ensuring thinking budget < max_tokens
// - Removing thinkingConfig if budget < minimum allowed
//
// Returns the normalized budget and updated payload.
// Returns budget=-2 as a sentinel indicating thinkingConfig was removed entirely.
func (a *Applier) normalizeClaudeBudget(budget int, payload []byte, modelInfo *registry.ModelInfo) (int, []byte) {
if modelInfo == nil {
return budget, payload
}
// Get effective max tokens
effectiveMax, setDefaultMax := a.effectiveMaxTokens(payload, modelInfo)
if effectiveMax > 0 && budget >= effectiveMax {
budget = effectiveMax - 1
}
// Check minimum budget
minBudget := 0
if modelInfo.Thinking != nil {
minBudget = modelInfo.Thinking.Min
}
if minBudget > 0 && budget >= 0 && budget < minBudget {
// Budget is below minimum, remove thinking config entirely
payload, _ = sjson.DeleteBytes(payload, "request.generationConfig.thinkingConfig")
return -2, payload
}
// Set default max tokens if needed
if setDefaultMax && effectiveMax > 0 {
payload, _ = sjson.SetBytes(payload, "request.generationConfig.maxOutputTokens", effectiveMax)
}
return budget, payload
}
// effectiveMaxTokens returns the max tokens to cap thinking:
// prefer request-provided maxOutputTokens; otherwise fall back to model default.
// The boolean indicates whether the value came from the model default (and thus should be written back).
func (a *Applier) effectiveMaxTokens(payload []byte, modelInfo *registry.ModelInfo) (max int, fromModel bool) {
if maxTok := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxTok.Exists() && maxTok.Int() > 0 {
return int(maxTok.Int()), false
}
if modelInfo != nil && modelInfo.MaxCompletionTokens > 0 {
return modelInfo.MaxCompletionTokens, true
}
return 0, false
}

View File

@@ -22,9 +22,7 @@ func NewApplier() *Applier {
} }
func init() { func init() {
applier := NewApplier() thinking.RegisterProvider("gemini-cli", NewApplier())
thinking.RegisterProvider("gemini-cli", applier)
thinking.RegisterProvider("antigravity", applier)
} }
// Apply applies thinking configuration to Gemini CLI request body. // Apply applies thinking configuration to Gemini CLI request body.

View File

@@ -8,6 +8,7 @@ import (
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator"
// Import provider packages to trigger init() registration of ProviderAppliers // Import provider packages to trigger init() registration of ProviderAppliers
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/antigravity"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/claude" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/claude"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/codex" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/codex"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/gemini" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/gemini"