mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-03 13:00:52 +08:00
refactor: improve thinking logic
This commit is contained in:
hkfires
116
internal/thinking/provider/claude/apply.go
Normal file
116
internal/thinking/provider/claude/apply.go
Normal file
@@ -0,0 +1,116 @@
|
||||
// Package claude implements thinking configuration scaffolding for Claude models.
|
||||
//
|
||||
// Claude models use the thinking.budget_tokens format with values in the range
|
||||
// 1024-128000. Some Claude models support ZeroAllowed (sonnet-4-5, opus-4-5),
|
||||
// while older models do not.
|
||||
// See: _bmad-output/planning-artifacts/architecture.md#Epic-6
|
||||
package claude
|
||||
|
||||
import (
|
||||
"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 implements thinking.ProviderApplier for Claude models.
|
||||
// This applier is stateless and holds no configuration.
|
||||
type Applier struct{}
|
||||
|
||||
// NewApplier creates a new Claude thinking applier.
|
||||
func NewApplier() *Applier {
|
||||
return &Applier{}
|
||||
}
|
||||
|
||||
func init() {
|
||||
thinking.RegisterProvider("claude", NewApplier())
|
||||
}
|
||||
|
||||
// Apply applies thinking configuration to Claude request body.
|
||||
//
|
||||
// IMPORTANT: This method expects config to be pre-validated by thinking.ValidateConfig.
|
||||
// ValidateConfig handles:
|
||||
// - Mode conversion (Level→Budget, Auto→Budget)
|
||||
// - Budget clamping to model range
|
||||
// - ZeroAllowed constraint enforcement
|
||||
//
|
||||
// Apply only processes ModeBudget and ModeNone; other modes are passed through unchanged.
|
||||
//
|
||||
// Expected output format when enabled:
|
||||
//
|
||||
// {
|
||||
// "thinking": {
|
||||
// "type": "enabled",
|
||||
// "budget_tokens": 16384
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// Expected output format when disabled:
|
||||
//
|
||||
// {
|
||||
// "thinking": {
|
||||
// "type": "disabled"
|
||||
// }
|
||||
// }
|
||||
func (a *Applier) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) {
|
||||
if modelInfo == nil {
|
||||
return body, nil
|
||||
}
|
||||
if modelInfo.Thinking == nil {
|
||||
if modelInfo.Type == "" {
|
||||
modelID := modelInfo.ID
|
||||
if modelID == "" {
|
||||
modelID = "unknown"
|
||||
}
|
||||
return nil, thinking.NewThinkingErrorWithModel(thinking.ErrThinkingNotSupported, "thinking not supported for this model", modelID)
|
||||
}
|
||||
return applyCompatibleClaude(body, config)
|
||||
}
|
||||
|
||||
// Only process ModeBudget and ModeNone; other modes pass through
|
||||
// (caller should use ValidateConfig first to normalize modes)
|
||||
if config.Mode != thinking.ModeBudget && config.Mode != thinking.ModeNone {
|
||||
return body, nil
|
||||
}
|
||||
|
||||
if len(body) == 0 || !gjson.ValidBytes(body) {
|
||||
body = []byte(`{}`)
|
||||
}
|
||||
|
||||
// Budget is expected to be pre-validated by ValidateConfig (clamped, ZeroAllowed enforced)
|
||||
// Decide enabled/disabled based on budget value
|
||||
if config.Budget == 0 {
|
||||
result, _ := sjson.SetBytes(body, "thinking.type", "disabled")
|
||||
result, _ = sjson.DeleteBytes(result, "thinking.budget_tokens")
|
||||
return result, nil
|
||||
}
|
||||
|
||||
result, _ := sjson.SetBytes(body, "thinking.type", "enabled")
|
||||
result, _ = sjson.SetBytes(result, "thinking.budget_tokens", config.Budget)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func applyCompatibleClaude(body []byte, config thinking.ThinkingConfig) ([]byte, error) {
|
||||
if config.Mode != thinking.ModeBudget && config.Mode != thinking.ModeNone && config.Mode != thinking.ModeAuto {
|
||||
return body, nil
|
||||
}
|
||||
|
||||
if len(body) == 0 || !gjson.ValidBytes(body) {
|
||||
body = []byte(`{}`)
|
||||
}
|
||||
|
||||
switch config.Mode {
|
||||
case thinking.ModeNone:
|
||||
result, _ := sjson.SetBytes(body, "thinking.type", "disabled")
|
||||
result, _ = sjson.DeleteBytes(result, "thinking.budget_tokens")
|
||||
return result, nil
|
||||
case thinking.ModeAuto:
|
||||
result, _ := sjson.SetBytes(body, "thinking.type", "enabled")
|
||||
result, _ = sjson.DeleteBytes(result, "thinking.budget_tokens")
|
||||
return result, nil
|
||||
default:
|
||||
result, _ := sjson.SetBytes(body, "thinking.type", "enabled")
|
||||
result, _ = sjson.SetBytes(result, "thinking.budget_tokens", config.Budget)
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
288
internal/thinking/provider/claude/apply_test.go
Normal file
288
internal/thinking/provider/claude/apply_test.go
Normal file
@@ -0,0 +1,288 @@
|
||||
// Package claude implements thinking configuration for Claude models.
|
||||
package claude
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// Unit Tests: Applier Creation and Interface
|
||||
// =============================================================================
|
||||
|
||||
func TestNewApplier(t *testing.T) {
|
||||
applier := NewApplier()
|
||||
if applier == nil {
|
||||
t.Fatal("NewApplier() returned nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplierImplementsInterface(t *testing.T) {
|
||||
var _ thinking.ProviderApplier = (*Applier)(nil)
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Unit Tests: Budget and Disable Logic (Pre-validated Config)
|
||||
// =============================================================================
|
||||
|
||||
// TestClaudeApplyBudgetAndNone tests budget values and disable modes.
|
||||
// NOTE: These tests assume config has been pre-validated by ValidateConfig.
|
||||
// Apply trusts the input and does not perform clamping.
|
||||
func TestClaudeApplyBudgetAndNone(t *testing.T) {
|
||||
applier := NewApplier()
|
||||
modelInfo := buildClaudeModelInfo()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
config thinking.ThinkingConfig
|
||||
wantType string
|
||||
wantBudget int
|
||||
wantBudgetOK bool
|
||||
}{
|
||||
// Valid pre-validated budget values
|
||||
{"budget 16k", thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 16384}, "enabled", 16384, true},
|
||||
{"budget min", thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 1024}, "enabled", 1024, true},
|
||||
{"budget max", thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 128000}, "enabled", 128000, true},
|
||||
{"budget mid", thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 50000}, "enabled", 50000, true},
|
||||
// Disable cases
|
||||
{"budget zero disables", thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 0}, "disabled", 0, false},
|
||||
{"mode none disables", thinking.ThinkingConfig{Mode: thinking.ModeNone, Budget: 0}, "disabled", 0, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result, err := applier.Apply([]byte(`{}`), tt.config, modelInfo)
|
||||
if err != nil {
|
||||
t.Fatalf("Apply() error = %v", err)
|
||||
}
|
||||
|
||||
thinkingType := gjson.GetBytes(result, "thinking.type").String()
|
||||
if thinkingType != tt.wantType {
|
||||
t.Fatalf("thinking.type = %q, want %q", thinkingType, tt.wantType)
|
||||
}
|
||||
|
||||
budgetValue := gjson.GetBytes(result, "thinking.budget_tokens")
|
||||
if budgetValue.Exists() != tt.wantBudgetOK {
|
||||
t.Fatalf("thinking.budget_tokens exists = %v, want %v", budgetValue.Exists(), tt.wantBudgetOK)
|
||||
}
|
||||
if tt.wantBudgetOK {
|
||||
if got := int(budgetValue.Int()); got != tt.wantBudget {
|
||||
t.Fatalf("thinking.budget_tokens = %d, want %d", got, tt.wantBudget)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestClaudeApplyPassthroughBudget tests that Apply trusts pre-validated budget values.
|
||||
// It does NOT perform clamping - that's ValidateConfig's responsibility.
|
||||
func TestClaudeApplyPassthroughBudget(t *testing.T) {
|
||||
applier := NewApplier()
|
||||
modelInfo := buildClaudeModelInfo()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
config thinking.ThinkingConfig
|
||||
wantBudget int
|
||||
}{
|
||||
// Apply should pass through the budget value as-is
|
||||
// (ValidateConfig would have clamped these, but Apply trusts the input)
|
||||
{"passes through any budget", thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 500}, 500},
|
||||
{"passes through large budget", thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 200000}, 200000},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result, err := applier.Apply([]byte(`{}`), tt.config, modelInfo)
|
||||
if err != nil {
|
||||
t.Fatalf("Apply() error = %v", err)
|
||||
}
|
||||
|
||||
if got := int(gjson.GetBytes(result, "thinking.budget_tokens").Int()); got != tt.wantBudget {
|
||||
t.Fatalf("thinking.budget_tokens = %d, want %d (passthrough)", got, tt.wantBudget)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Unit Tests: Mode Passthrough (Strict Layering)
|
||||
// =============================================================================
|
||||
|
||||
// TestClaudeApplyModePassthrough tests that non-Budget/None modes pass through unchanged.
|
||||
// Apply expects ValidateConfig to have already converted Level/Auto to Budget.
|
||||
func TestClaudeApplyModePassthrough(t *testing.T) {
|
||||
applier := NewApplier()
|
||||
modelInfo := buildClaudeModelInfo()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
config thinking.ThinkingConfig
|
||||
body string
|
||||
}{
|
||||
{"ModeLevel passes through", thinking.ThinkingConfig{Mode: thinking.ModeLevel, Level: "high"}, `{"model":"test"}`},
|
||||
{"ModeAuto passes through", thinking.ThinkingConfig{Mode: thinking.ModeAuto, Budget: -1}, `{"model":"test"}`},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result, err := applier.Apply([]byte(tt.body), tt.config, modelInfo)
|
||||
if err != nil {
|
||||
t.Fatalf("Apply() error = %v", err)
|
||||
}
|
||||
|
||||
// Should return body unchanged
|
||||
if string(result) != tt.body {
|
||||
t.Fatalf("Apply() = %s, want %s (passthrough)", string(result), tt.body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Unit Tests: Output Format
|
||||
// =============================================================================
|
||||
|
||||
// TestClaudeApplyOutputFormat tests the exact JSON output format.
|
||||
//
|
||||
// Claude expects:
|
||||
//
|
||||
// {
|
||||
// "thinking": {
|
||||
// "type": "enabled",
|
||||
// "budget_tokens": 16384
|
||||
// }
|
||||
// }
|
||||
func TestClaudeApplyOutputFormat(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
config thinking.ThinkingConfig
|
||||
wantJSON string
|
||||
}{
|
||||
{
|
||||
"enabled with budget",
|
||||
thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 16384},
|
||||
`{"thinking":{"type":"enabled","budget_tokens":16384}}`,
|
||||
},
|
||||
{
|
||||
"disabled",
|
||||
thinking.ThinkingConfig{Mode: thinking.ModeNone, Budget: 0},
|
||||
`{"thinking":{"type":"disabled"}}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
applier := NewApplier()
|
||||
modelInfo := buildClaudeModelInfo()
|
||||
|
||||
result, err := applier.Apply([]byte(`{}`), tt.config, modelInfo)
|
||||
if err != nil {
|
||||
t.Fatalf("Apply() error = %v", err)
|
||||
}
|
||||
if string(result) != tt.wantJSON {
|
||||
t.Fatalf("Apply() = %s, want %s", result, tt.wantJSON)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Unit Tests: Body Merging
|
||||
// =============================================================================
|
||||
|
||||
// TestClaudeApplyWithExistingBody tests applying config to existing request body.
|
||||
func TestClaudeApplyWithExistingBody(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
config thinking.ThinkingConfig
|
||||
wantBody string
|
||||
}{
|
||||
{
|
||||
"add to empty body",
|
||||
`{}`,
|
||||
thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 16384},
|
||||
`{"thinking":{"type":"enabled","budget_tokens":16384}}`,
|
||||
},
|
||||
{
|
||||
"preserve existing fields",
|
||||
`{"model":"claude-sonnet-4-5","messages":[]}`,
|
||||
thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 8192},
|
||||
`{"model":"claude-sonnet-4-5","messages":[],"thinking":{"type":"enabled","budget_tokens":8192}}`,
|
||||
},
|
||||
{
|
||||
"override existing thinking",
|
||||
`{"thinking":{"type":"enabled","budget_tokens":1000}}`,
|
||||
thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 16384},
|
||||
`{"thinking":{"type":"enabled","budget_tokens":16384}}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
applier := NewApplier()
|
||||
modelInfo := buildClaudeModelInfo()
|
||||
|
||||
result, err := applier.Apply([]byte(tt.body), tt.config, modelInfo)
|
||||
if err != nil {
|
||||
t.Fatalf("Apply() error = %v", err)
|
||||
}
|
||||
if string(result) != tt.wantBody {
|
||||
t.Fatalf("Apply() = %s, want %s", result, tt.wantBody)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestClaudeApplyWithNilBody tests handling of nil/empty body.
|
||||
func TestClaudeApplyWithNilBody(t *testing.T) {
|
||||
applier := NewApplier()
|
||||
modelInfo := buildClaudeModelInfo()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
wantBudget int
|
||||
}{
|
||||
{"nil body", nil, 16384},
|
||||
{"empty body", []byte{}, 16384},
|
||||
{"empty object", []byte(`{}`), 16384},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
config := thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 16384}
|
||||
result, err := applier.Apply(tt.body, config, modelInfo)
|
||||
if err != nil {
|
||||
t.Fatalf("Apply() error = %v", err)
|
||||
}
|
||||
|
||||
if got := gjson.GetBytes(result, "thinking.type").String(); got != "enabled" {
|
||||
t.Fatalf("thinking.type = %q, want %q", got, "enabled")
|
||||
}
|
||||
if got := int(gjson.GetBytes(result, "thinking.budget_tokens").Int()); got != tt.wantBudget {
|
||||
t.Fatalf("thinking.budget_tokens = %d, want %d", got, tt.wantBudget)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Helper Functions
|
||||
// =============================================================================
|
||||
|
||||
func buildClaudeModelInfo() *registry.ModelInfo {
|
||||
return ®istry.ModelInfo{
|
||||
ID: "claude-sonnet-4-5",
|
||||
Thinking: ®istry.ThinkingSupport{
|
||||
Min: 1024,
|
||||
Max: 128000,
|
||||
ZeroAllowed: true,
|
||||
DynamicAllowed: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user