12d0f1f · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat: add reasoning effort support and new providers

12d0f1fKujtim Hoxha1 year ago

internal/config/config.go+15−25

compare with main

unifiedsplit

@@ -65,6 +65,7 @@ type Model struct {

65 DefaultMaxTokens int64 `json:"default_max_tokens"`

66 CanReason bool `json:"can_reason"`

67 ReasoningEffort string `json:"reasoning_effort"`

68+ HasReasoningEffort bool `json:"has_reasoning_effort"`

68 SupportsImages bool `json:"supports_attachments"`

69 SupportsImages bool `json:"supports_attachments"`

69}

70}

@@ -156,8 +157,9 @@ type Options struct {

156}

157}

157

158

158type PreferredModel struct {

159type PreferredModel struct {

159− ModelID string `json:"model_id"`

160+ ModelID string `json:"model_id"`

160− Provider provider.InferenceProvider `json:"provider"`

161+ Provider provider.InferenceProvider `json:"provider"`

162+ ReasoningEffort string `json:"reasoning_effort,omitempty"`

161}

163}

162

164

163type PreferredModels struct {

165type PreferredModels struct {

@@ -693,7 +695,7 @@ func defaultConfigBasedOnEnv() *Config {

693 }

695 }

694 providerConfig.BaseURL = baseURL

696 providerConfig.BaseURL = baseURL

695 for _, model := range p.Models {

697 for _, model := range p.Models {

696− providerConfig.Models = append(providerConfig.Models, Model{

698+ configModel := Model{

697 ID: model.ID,

699 ID: model.ID,

698 Name: model.Name,

700 Name: model.Name,

699 CostPer1MIn: model.CostPer1MIn,

701 CostPer1MIn: model.CostPer1MIn,

@@ -704,7 +706,13 @@ func defaultConfigBasedOnEnv() *Config {

704 DefaultMaxTokens: model.DefaultMaxTokens,

706 DefaultMaxTokens: model.DefaultMaxTokens,

705 CanReason: model.CanReason,

707 CanReason: model.CanReason,

706 SupportsImages: model.SupportsImages,

708 SupportsImages: model.SupportsImages,

707− })

709+ }

710+ // Set reasoning effort for reasoning models

711+ if model.HasReasoningEffort && model.DefaultReasoningEffort != "" {

712+ configModel.HasReasoningEffort = model.HasReasoningEffort

713+ configModel.ReasoningEffort = model.DefaultReasoningEffort

714+ }

715+ providerConfig.Models = append(providerConfig.Models, configModel)

708 }

716 }

709 cfg.Providers[p.ID] = providerConfig

717 cfg.Providers[p.ID] = providerConfig

710 }

718 }

@@ -980,25 +988,13 @@ func (c *Config) validateProviders(errors *ValidationErrors) {

980 }

988 }

981

989

982 // Validate provider type

990 // Validate provider type

983− validType := false

984− for _, vt := range validTypes {

985− if providerConfig.ProviderType == vt {

991+ validType := slices.Contains(validTypes, providerConfig.ProviderType)

986− validType = true

987− break

988− }

989− }

990 if !validType {

992 if !validType {

991 errors.Add(fieldPrefix+".provider_type", fmt.Sprintf("invalid provider type: %s", providerConfig.ProviderType))

993 errors.Add(fieldPrefix+".provider_type", fmt.Sprintf("invalid provider type: %s", providerConfig.ProviderType))

992 }

994 }

993

995

994 // Validate custom providers

996 // Validate custom providers

995− isKnownProvider := false

997+ isKnownProvider := slices.Contains(knownProviders, providerID)

996− for _, kp := range knownProviders {

997− if providerID == kp {

998− isKnownProvider = true

999− break

1000− }

1001− }

1002

998

1003 if !isKnownProvider {

999 if !isKnownProvider {

1004 // Custom provider validation

1000 // Custom provider validation

@@ -1200,13 +1196,7 @@ func (c *Config) validateAgents(errors *ValidationErrors) {

1200 // Validate allowed tools

1196 // Validate allowed tools

1201 if agent.AllowedTools != nil {

1197 if agent.AllowedTools != nil {

1202 for i, tool := range agent.AllowedTools {

1198 for i, tool := range agent.AllowedTools {

1203− validTool := false

1199+ validTool := slices.Contains(validTools, tool)

1204− for _, vt := range validTools {

1205− if tool == vt {

1206− validTool = true

1207− break

1208− }

1209− }

1210 if !validTool {

1200 if !validTool {

1211 errors.Add(fmt.Sprintf("%s.allowed_tools[%d]", fieldPrefix, i), fmt.Sprintf("unknown tool: %s", tool))

1201 errors.Add(fmt.Sprintf("%s.allowed_tools[%d]", fieldPrefix, i), fmt.Sprintf("unknown tool: %s", tool))

1212 }

1202 }

internal/config/provider_mock.go+80

@@ -136,6 +136,34 @@ func MockProviders() []provider.Provider {

136 CanReason: false,

137 SupportsImages: true,

138 },

139+ {

140+ ID: "o1-preview",

141+ Name: "o1-preview",

142+ CostPer1MIn: 15.0,

143+ CostPer1MOut: 60.0,

144+ CostPer1MInCached: 0.0,

145+ CostPer1MOutCached: 0.0,

146+ ContextWindow: 128000,

147+ DefaultMaxTokens: 32768,

148+ CanReason: true,

149+ HasReasoningEffort: true,

150+ DefaultReasoningEffort: "medium",

151+ SupportsImages: true,

152+ },

153+ {

154+ ID: "o1-mini",

155+ Name: "o1-mini",

156+ CostPer1MIn: 3.0,

157+ CostPer1MOut: 12.0,

158+ CostPer1MInCached: 0.0,

159+ CostPer1MOutCached: 0.0,

160+ ContextWindow: 128000,

161+ DefaultMaxTokens: 65536,

162+ CanReason: true,

163+ HasReasoningEffort: true,

164+ DefaultReasoningEffort: "medium",

165+ SupportsImages: true,

166+ },

139 },

167 },

140 },

168 },

141 {

169 {

@@ -173,5 +201,57 @@ func MockProviders() []provider.Provider {

173 },

201 },

174 },

202 },

175 },

203 },

204+ {

205+ Name: "xAI",

206+ ID: provider.InferenceProviderXAI,

207+ APIKey: "$XAI_API_KEY",

208+ APIEndpoint: "https://api.x.ai/v1",

209+ Type: provider.TypeXAI,

210+ DefaultLargeModelID: "grok-beta",

211+ DefaultSmallModelID: "grok-beta",

212+ Models: []provider.Model{

213+ {

214+ ID: "grok-beta",

215+ Name: "Grok Beta",

216+ CostPer1MIn: 5.0,

217+ CostPer1MOut: 15.0,

218+ ContextWindow: 131072,

219+ DefaultMaxTokens: 4096,

220+ CanReason: false,

221+ SupportsImages: true,

222+ },

223+ },

224+ },

225+ {

226+ Name: "OpenRouter",

227+ ID: provider.InferenceProviderOpenRouter,

228+ APIKey: "$OPENROUTER_API_KEY",

229+ APIEndpoint: "https://openrouter.ai/api/v1",

230+ Type: provider.TypeOpenAI,

231+ DefaultLargeModelID: "anthropic/claude-3.5-sonnet",

232+ DefaultSmallModelID: "anthropic/claude-3.5-haiku",

233+ Models: []provider.Model{

234+ {

235+ ID: "anthropic/claude-3.5-sonnet",

236+ Name: "Claude 3.5 Sonnet",

237+ CostPer1MIn: 3.0,

238+ CostPer1MOut: 15.0,

239+ ContextWindow: 200000,

240+ DefaultMaxTokens: 8192,

241+ CanReason: false,

242+ SupportsImages: true,

243+ },

244+ {

245+ ID: "anthropic/claude-3.5-haiku",

246+ Name: "Claude 3.5 Haiku",

247+ CostPer1MIn: 0.8,

248+ CostPer1MOut: 4.0,

249+ ContextWindow: 200000,

250+ DefaultMaxTokens: 8192,

251+ CanReason: false,

252+ SupportsImages: true,

253+ },

254+ },

255+ },

176 }

256 }

177}

257}

internal/config/provider_test.go+175

@@ -1,6 +1,7 @@

1package config

3import (

4+ "encoding/json"

4 "testing"

5 "testing"

6 "github.com/charmbracelet/crush/internal/fur/provider"

7 "github.com/charmbracelet/crush/internal/fur/provider"

@@ -103,3 +104,177 @@ func TestResetProviders(t *testing.T) {

103 // Should get the same mock data

104 // Should get the same mock data

104 assert.Equal(t, len(providers1), len(providers2))

105 assert.Equal(t, len(providers1), len(providers2))

105}

106}

107+

108+func TestReasoningEffortSupport(t *testing.T) {

109+ originalUseMock := UseMockProviders

110+ UseMockProviders = true

111+ defer func() {

112+ UseMockProviders = originalUseMock

113+ ResetProviders()

114+ }()

115+

116+ ResetProviders()

117+ providers := Providers()

118+

119+ var openaiProvider provider.Provider

120+ for _, p := range providers {

121+ if p.ID == provider.InferenceProviderOpenAI {

122+ openaiProvider = p

123+ break

124+ }

125+ }

126+ require.NotEmpty(t, openaiProvider.ID)

127+

128+ var reasoningModel, nonReasoningModel provider.Model

129+ for _, model := range openaiProvider.Models {

130+ if model.CanReason && model.HasReasoningEffort {

131+ reasoningModel = model

132+ } else if !model.CanReason {

133+ nonReasoningModel = model

134+ }

135+ }

136+

137+ require.NotEmpty(t, reasoningModel.ID)

138+ assert.Equal(t, "medium", reasoningModel.DefaultReasoningEffort)

139+ assert.True(t, reasoningModel.HasReasoningEffort)

140+

141+ require.NotEmpty(t, nonReasoningModel.ID)

142+ assert.False(t, nonReasoningModel.HasReasoningEffort)

143+ assert.Empty(t, nonReasoningModel.DefaultReasoningEffort)

144+}

145+

146+func TestReasoningEffortConfigTransfer(t *testing.T) {

147+ originalUseMock := UseMockProviders

148+ UseMockProviders = true

149+ defer func() {

150+ UseMockProviders = originalUseMock

151+ ResetProviders()

152+ }()

153+

154+ ResetProviders()

155+ t.Setenv("OPENAI_API_KEY", "test-openai-key")

156+

157+ cfg, err := Init(t.TempDir(), false)

158+ require.NoError(t, err)

159+

160+ openaiProviderConfig, exists := cfg.Providers[provider.InferenceProviderOpenAI]

161+ require.True(t, exists)

162+

163+ var foundReasoning, foundNonReasoning bool

164+ for _, model := range openaiProviderConfig.Models {

165+ if model.CanReason && model.HasReasoningEffort && model.ReasoningEffort != "" {

166+ assert.Equal(t, "medium", model.ReasoningEffort)

167+ assert.True(t, model.HasReasoningEffort)

168+ foundReasoning = true

169+ } else if !model.CanReason {

170+ assert.Empty(t, model.ReasoningEffort)

171+ assert.False(t, model.HasReasoningEffort)

172+ foundNonReasoning = true

173+ }

174+ }

175+

176+ assert.True(t, foundReasoning, "Should find at least one reasoning model")

177+ assert.True(t, foundNonReasoning, "Should find at least one non-reasoning model")

178+}

179+

180+func TestNewProviders(t *testing.T) {

181+ originalUseMock := UseMockProviders

182+ UseMockProviders = true

183+ defer func() {

184+ UseMockProviders = originalUseMock

185+ ResetProviders()

186+ }()

187+

188+ ResetProviders()

189+ providers := Providers()

190+ require.NotEmpty(t, providers)

191+

192+ var xaiProvider, openRouterProvider provider.Provider

193+ for _, p := range providers {

194+ switch p.ID {

195+ case provider.InferenceProviderXAI:

196+ xaiProvider = p

197+ case provider.InferenceProviderOpenRouter:

198+ openRouterProvider = p

199+ }

200+ }

201+

202+ require.NotEmpty(t, xaiProvider.ID)

203+ assert.Equal(t, "xAI", xaiProvider.Name)

204+ assert.Equal(t, "grok-beta", xaiProvider.DefaultLargeModelID)

205+

206+ require.NotEmpty(t, openRouterProvider.ID)

207+ assert.Equal(t, "OpenRouter", openRouterProvider.Name)

208+ assert.Equal(t, "anthropic/claude-3.5-sonnet", openRouterProvider.DefaultLargeModelID)

209+}

210+

211+func TestO1ModelsInMockProvider(t *testing.T) {

212+ originalUseMock := UseMockProviders

213+ UseMockProviders = true

214+ defer func() {

215+ UseMockProviders = originalUseMock

216+ ResetProviders()

217+ }()

218+

219+ ResetProviders()

220+ providers := Providers()

221+

222+ var openaiProvider provider.Provider

223+ for _, p := range providers {

224+ if p.ID == provider.InferenceProviderOpenAI {

225+ openaiProvider = p

226+ break

227+ }

228+ }

229+ require.NotEmpty(t, openaiProvider.ID)

230+

231+ modelTests := []struct {

232+ id string

233+ name string

234+ }{

235+ {"o1-preview", "o1-preview"},

236+ {"o1-mini", "o1-mini"},

237+ }

238+

239+ for _, test := range modelTests {

240+ var model provider.Model

241+ var found bool

242+ for _, m := range openaiProvider.Models {

243+ if m.ID == test.id {

244+ model = m

245+ found = true

246+ break

247+ }

248+ }

249+ require.True(t, found, "Should find %s model", test.id)

250+ assert.Equal(t, test.name, model.Name)

251+ assert.True(t, model.CanReason)

252+ assert.True(t, model.HasReasoningEffort)

253+ assert.Equal(t, "medium", model.DefaultReasoningEffort)

254+ }

255+}

256+

257+func TestPreferredModelReasoningEffort(t *testing.T) {

258+ // Test that PreferredModel struct can hold reasoning effort

259+ preferredModel := PreferredModel{

260+ ModelID: "o1-preview",

261+ Provider: provider.InferenceProviderOpenAI,

262+ ReasoningEffort: "high",

263+ }

264+

265+ assert.Equal(t, "o1-preview", preferredModel.ModelID)

266+ assert.Equal(t, provider.InferenceProviderOpenAI, preferredModel.Provider)

267+ assert.Equal(t, "high", preferredModel.ReasoningEffort)

268+

269+ // Test JSON marshaling/unmarshaling

270+ jsonData, err := json.Marshal(preferredModel)

271+ require.NoError(t, err)

272+

273+ var unmarshaled PreferredModel

274+ err = json.Unmarshal(jsonData, &unmarshaled)

275+ require.NoError(t, err)

276+

277+ assert.Equal(t, preferredModel.ModelID, unmarshaled.ModelID)

278+ assert.Equal(t, preferredModel.Provider, unmarshaled.Provider)

279+ assert.Equal(t, preferredModel.ReasoningEffort, unmarshaled.ReasoningEffort)

280+}

internal/llm/provider/gemini.go−21

@@ -407,27 +407,6 @@ func (g *geminiClient) shouldRetry(attempts int, err error) (bool, int64, error)

407 return true, int64(retryMs), nil

408}

409

410−func (g *geminiClient) toolCalls(resp *genai.GenerateContentResponse) []message.ToolCall {

411− var toolCalls []message.ToolCall

412−

413− if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {

414− for _, part := range resp.Candidates[0].Content.Parts {

415− if part.FunctionCall != nil {

416− id := "call_" + uuid.New().String()

417− args, _ := json.Marshal(part.FunctionCall.Args)

418− toolCalls = append(toolCalls, message.ToolCall{

419− ID: id,

420− Name: part.FunctionCall.Name,

421− Input: string(args),

422− Type: "function",

423− })

424− }

425− }

426− }

427−

428− return toolCalls

429−}

430−

431func (g *geminiClient) usage(resp *genai.GenerateContentResponse) TokenUsage {

410func (g *geminiClient) usage(resp *genai.GenerateContentResponse) TokenUsage {

432 if resp == nil || resp.UsageMetadata == nil {

411 if resp == nil || resp.UsageMetadata == nil {

433 return TokenUsage{}

412 return TokenUsage{}

internal/llm/provider/openai.go+13−11

@@ -18,23 +18,14 @@ import (

18 "github.com/openai/openai-go/shared"

19)

21−type openaiOptions struct {

22− reasoningEffort string

23−}

24−

25type openaiClient struct {

21type openaiClient struct {

26 providerOptions providerClientOptions

22 providerOptions providerClientOptions

27− options openaiOptions

28 client openai.Client

23 client openai.Client

29}

24}

31type OpenAIClient ProviderClient

26type OpenAIClient ProviderClient

33func newOpenAIClient(opts providerClientOptions) OpenAIClient {

28func newOpenAIClient(opts providerClientOptions) OpenAIClient {

34− openaiOpts := openaiOptions{

35− reasoningEffort: "medium",

36− }

37−

38 openaiClientOptions := []option.RequestOption{}

29 openaiClientOptions := []option.RequestOption{}

39 if opts.apiKey != "" {

30 if opts.apiKey != "" {

40 openaiClientOptions = append(openaiClientOptions, option.WithAPIKey(opts.apiKey))

31 openaiClientOptions = append(openaiClientOptions, option.WithAPIKey(opts.apiKey))

@@ -52,7 +43,6 @@ func newOpenAIClient(opts providerClientOptions) OpenAIClient {

52 client := openai.NewClient(openaiClientOptions...)

43 client := openai.NewClient(openaiClientOptions...)

53 return &openaiClient{

44 return &openaiClient{

54 providerOptions: opts,

45 providerOptions: opts,

55− options: openaiOpts,

56 client: client,

46 client: client,

57 }

47 }

58}

48}

@@ -153,6 +143,18 @@ func (o *openaiClient) finishReason(reason string) message.FinishReason {

153

143

154func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessageParamUnion, tools []openai.ChatCompletionToolParam) openai.ChatCompletionNewParams {

144func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessageParamUnion, tools []openai.ChatCompletionToolParam) openai.ChatCompletionNewParams {

155 model := o.providerOptions.model(o.providerOptions.modelType)

145 model := o.providerOptions.model(o.providerOptions.modelType)

146+ cfg := config.Get()

147+

148+ modelConfig := cfg.Models.Large

149+ if o.providerOptions.modelType == config.SmallModel {

150+ modelConfig = cfg.Models.Small

151+ }

152+

153+ reasoningEffort := model.ReasoningEffort

154+ if modelConfig.ReasoningEffort != "" {

155+ reasoningEffort = modelConfig.ReasoningEffort

156+ }

157+

156 params := openai.ChatCompletionNewParams{

158 params := openai.ChatCompletionNewParams{

157 Model: openai.ChatModel(model.ID),

159 Model: openai.ChatModel(model.ID),

158 Messages: messages,

160 Messages: messages,

@@ -160,7 +162,7 @@ func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessagePar

160 }

162 }

161 if model.CanReason {

163 if model.CanReason {

162 params.MaxCompletionTokens = openai.Int(o.providerOptions.maxTokens)

164 params.MaxCompletionTokens = openai.Int(o.providerOptions.maxTokens)

163− switch o.options.reasoningEffort {

165+ switch reasoningEffort {

164 case "low":

166 case "low":

165 params.ReasoningEffort = shared.ReasoningEffortLow

167 params.ReasoningEffort = shared.ReasoningEffortLow

166 case "medium":

168 case "medium":