eb3550a · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

chore: change how max tokens works

eb3550aKujtim Hoxha1 year ago

internal/config/config.go+6−3

compare with main

unifiedsplit

@@ -157,9 +157,12 @@ type Options struct {

157}

158

159type PreferredModel struct {

160− ModelID string `json:"model_id"`

160+ ModelID string `json:"model_id"`

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

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

162+ // Overrides the default reasoning effort for this model

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

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

164+ // Overrides the default max tokens for this model

165+ MaxTokens int64 `json:"max_tokens,omitempty"`

163}

166}

164

167

165type PreferredModels struct {

168type PreferredModels struct {

internal/llm/agent/agent.go−6

@@ -147,7 +147,6 @@ func NewAgent(

147 opts := []provider.ProviderClientOption{

148 provider.WithModel(agentCfg.Model),

149 provider.WithSystemMessage(prompt.GetPrompt(promptID, providerCfg.ID)),

150− provider.WithMaxTokens(model.DefaultMaxTokens),

151 }

150 }

152 agentProvider, err := provider.NewProviderV2(providerCfg, opts...)

151 agentProvider, err := provider.NewProviderV2(providerCfg, opts...)

153 if err != nil {

152 if err != nil {

@@ -184,7 +183,6 @@ func NewAgent(

184 titleOpts := []provider.ProviderClientOption{

183 titleOpts := []provider.ProviderClientOption{

185 provider.WithModel(config.SmallModel),

184 provider.WithModel(config.SmallModel),

186 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),

185 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),

187− provider.WithMaxTokens(40),

188 }

186 }

189 titleProvider, err := provider.NewProviderV2(smallModelProviderCfg, titleOpts...)

187 titleProvider, err := provider.NewProviderV2(smallModelProviderCfg, titleOpts...)

190 if err != nil {

188 if err != nil {

@@ -193,7 +191,6 @@ func NewAgent(

193 summarizeOpts := []provider.ProviderClientOption{

191 summarizeOpts := []provider.ProviderClientOption{

194 provider.WithModel(config.SmallModel),

192 provider.WithModel(config.SmallModel),

195 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),

193 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),

196− provider.WithMaxTokens(smallModel.DefaultMaxTokens),

197 }

194 }

198 summarizeProvider, err := provider.NewProviderV2(smallModelProviderCfg, summarizeOpts...)

195 summarizeProvider, err := provider.NewProviderV2(smallModelProviderCfg, summarizeOpts...)

199 if err != nil {

196 if err != nil {

@@ -832,7 +829,6 @@ func (a *agent) UpdateModel() error {

832 opts := []provider.ProviderClientOption{

829 opts := []provider.ProviderClientOption{

833 provider.WithModel(a.agentCfg.Model),

830 provider.WithModel(a.agentCfg.Model),

834 provider.WithSystemMessage(prompt.GetPrompt(promptID, currentProviderCfg.ID)),

831 provider.WithSystemMessage(prompt.GetPrompt(promptID, currentProviderCfg.ID)),

835− provider.WithMaxTokens(model.DefaultMaxTokens),

836 }

832 }

837

833

838 newProvider, err := provider.NewProviderV2(currentProviderCfg, opts...)

834 newProvider, err := provider.NewProviderV2(currentProviderCfg, opts...)

@@ -877,7 +873,6 @@ func (a *agent) UpdateModel() error {

877 titleOpts := []provider.ProviderClientOption{

873 titleOpts := []provider.ProviderClientOption{

878 provider.WithModel(config.SmallModel),

874 provider.WithModel(config.SmallModel),

879 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),

875 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),

880− provider.WithMaxTokens(40),

881 }

876 }

882 newTitleProvider, err := provider.NewProviderV2(smallModelProviderCfg, titleOpts...)

877 newTitleProvider, err := provider.NewProviderV2(smallModelProviderCfg, titleOpts...)

883 if err != nil {

878 if err != nil {

@@ -888,7 +883,6 @@ func (a *agent) UpdateModel() error {

888 summarizeOpts := []provider.ProviderClientOption{

883 summarizeOpts := []provider.ProviderClientOption{

889 provider.WithModel(config.SmallModel),

884 provider.WithModel(config.SmallModel),

890 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),

885 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),

891− provider.WithMaxTokens(smallModel.DefaultMaxTokens),

892 }

886 }

893 newSummarizeProvider, err := provider.NewProviderV2(smallModelProviderCfg, summarizeOpts...)

887 newSummarizeProvider, err := provider.NewProviderV2(smallModelProviderCfg, summarizeOpts...)

894 if err != nil {

888 if err != nil {

internal/llm/provider/anthropic.go+11−1

@@ -164,9 +164,19 @@ func (a *anthropicClient) preparedMessages(messages []anthropic.MessageParam, to

164 // }

165 // }

166

167+ cfg := config.Get()

168+ modelConfig := cfg.Models.Large

169+ if a.providerOptions.modelType == config.SmallModel {

170+ modelConfig = cfg.Models.Small

171+ }

172+ maxTokens := model.DefaultMaxTokens

173+ if modelConfig.MaxTokens > 0 {

174+ maxTokens = modelConfig.MaxTokens

175+ }

176+

167 return anthropic.MessageNewParams{

177 return anthropic.MessageNewParams{

168 Model: anthropic.Model(model.ID),

178 Model: anthropic.Model(model.ID),

169− MaxTokens: a.providerOptions.maxTokens,

179+ MaxTokens: maxTokens,

170 Temperature: temperature,

180 Temperature: temperature,

171 Messages: messages,

181 Messages: messages,

172 Tools: tools,

182 Tools: tools,

internal/llm/provider/gemini.go+21−5

@@ -155,17 +155,26 @@ func (g *geminiClient) finishReason(reason genai.FinishReason) message.FinishRea

155func (g *geminiClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {

156 // Convert messages

157 geminiMessages := g.convertMessages(messages)

158−

158+ model := g.providerOptions.model(g.providerOptions.modelType)

159 cfg := config.Get()

160 if cfg.Options.Debug {

161 jsonData, _ := json.Marshal(geminiMessages)

162 logging.Debug("Prepared messages", "messages", string(jsonData))

163 }

164

165+ modelConfig := cfg.Models.Large

166+ if g.providerOptions.modelType == config.SmallModel {

167+ modelConfig = cfg.Models.Small

168+ }

169+

170+ maxTokens := model.DefaultMaxTokens

171+ if modelConfig.MaxTokens > 0 {

172+ maxTokens = modelConfig.MaxTokens

173+ }

165 history := geminiMessages[:len(geminiMessages)-1] // All but last message

174 history := geminiMessages[:len(geminiMessages)-1] // All but last message

166 lastMsg := geminiMessages[len(geminiMessages)-1]

175 lastMsg := geminiMessages[len(geminiMessages)-1]

167 config := &genai.GenerateContentConfig{

176 config := &genai.GenerateContentConfig{

168− MaxOutputTokens: int32(g.providerOptions.maxTokens),

177+ MaxOutputTokens: int32(maxTokens),

169 SystemInstruction: &genai.Content{

178 SystemInstruction: &genai.Content{

170 Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},

179 Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},

171 },

180 },

@@ -173,7 +182,6 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too

173 if len(tools) > 0 {

182 if len(tools) > 0 {

174 config.Tools = g.convertTools(tools)

183 config.Tools = g.convertTools(tools)

175 }

184 }

176− model := g.providerOptions.model(g.providerOptions.modelType)

177 chat, _ := g.client.Chats.Create(ctx, model.ID, config, history)

185 chat, _ := g.client.Chats.Create(ctx, model.ID, config, history)

178

186

179 attempts := 0

187 attempts := 0

@@ -245,16 +253,25 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t

245 // Convert messages

253 // Convert messages

246 geminiMessages := g.convertMessages(messages)

254 geminiMessages := g.convertMessages(messages)

247

255

256+ model := g.providerOptions.model(g.providerOptions.modelType)

248 cfg := config.Get()

257 cfg := config.Get()

249 if cfg.Options.Debug {

258 if cfg.Options.Debug {

250 jsonData, _ := json.Marshal(geminiMessages)

259 jsonData, _ := json.Marshal(geminiMessages)

251 logging.Debug("Prepared messages", "messages", string(jsonData))

260 logging.Debug("Prepared messages", "messages", string(jsonData))

252 }

261 }

253

262

263+ modelConfig := cfg.Models.Large

264+ if g.providerOptions.modelType == config.SmallModel {

265+ modelConfig = cfg.Models.Small

266+ }

267+ maxTokens := model.DefaultMaxTokens

268+ if modelConfig.MaxTokens > 0 {

269+ maxTokens = modelConfig.MaxTokens

270+ }

254 history := geminiMessages[:len(geminiMessages)-1] // All but last message

271 history := geminiMessages[:len(geminiMessages)-1] // All but last message

255 lastMsg := geminiMessages[len(geminiMessages)-1]

272 lastMsg := geminiMessages[len(geminiMessages)-1]

256 config := &genai.GenerateContentConfig{

273 config := &genai.GenerateContentConfig{

257− MaxOutputTokens: int32(g.providerOptions.maxTokens),

274+ MaxOutputTokens: int32(maxTokens),

258 SystemInstruction: &genai.Content{

275 SystemInstruction: &genai.Content{

259 Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},

276 Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},

260 },

277 },

@@ -262,7 +279,6 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t

262 if len(tools) > 0 {

279 if len(tools) > 0 {

263 config.Tools = g.convertTools(tools)

280 config.Tools = g.convertTools(tools)

264 }

281 }

265− model := g.providerOptions.model(g.providerOptions.modelType)

266 chat, _ := g.client.Chats.Create(ctx, model.ID, config, history)

282 chat, _ := g.client.Chats.Create(ctx, model.ID, config, history)

267

283

268 attempts := 0

284 attempts := 0

internal/llm/provider/openai.go+7−2

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

160 Messages: messages,

161 Tools: tools,

162 }

163+

164+ maxTokens := model.DefaultMaxTokens

165+ if modelConfig.MaxTokens > 0 {

166+ maxTokens = modelConfig.MaxTokens

167+ }

163 if model.CanReason {

168 if model.CanReason {

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

169+ params.MaxCompletionTokens = openai.Int(maxTokens)

165 switch reasoningEffort {

170 switch reasoningEffort {

166 case "low":

171 case "low":

167 params.ReasoningEffort = shared.ReasoningEffortLow

172 params.ReasoningEffort = shared.ReasoningEffortLow

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

173 params.ReasoningEffort = shared.ReasoningEffortMedium

178 params.ReasoningEffort = shared.ReasoningEffortMedium

174 }

179 }

175 } else {

180 } else {

176− params.MaxTokens = openai.Int(o.providerOptions.maxTokens)

181+ params.MaxTokens = openai.Int(maxTokens)

177 }

182 }

178

183

179 return params

184 return params

internal/llm/provider/provider.go−7

@@ -64,7 +64,6 @@ type providerClientOptions struct {

64 modelType config.ModelType

65 model func(config.ModelType) config.Model

66 disableCache bool

67− maxTokens int64

68 systemMessage string

67 systemMessage string

69 extraHeaders map[string]string

68 extraHeaders map[string]string

70 extraParams map[string]string

69 extraParams map[string]string

@@ -121,12 +120,6 @@ func WithDisableCache(disableCache bool) ProviderClientOption {

121 }

120 }

122}

121}

123

122

124−func WithMaxTokens(maxTokens int64) ProviderClientOption {

125− return func(options *providerClientOptions) {

126− options.maxTokens = maxTokens

127− }

128−}

129−

130func WithSystemMessage(systemMessage string) ProviderClientOption {

123func WithSystemMessage(systemMessage string) ProviderClientOption {

131 return func(options *providerClientOptions) {

124 return func(options *providerClientOptions) {

132 options.systemMessage = systemMessage

125 options.systemMessage = systemMessage