git.iko.soy/crush
on main help sign in
chore: change how max tokens works
commiteb3550a2bb26f9e4199a6339a3144dcd898dacbe authorKujtim Hoxha <kujtimii.h@gmail.com> dateSat 28 Jun 2025, 11:56 (1 year ago) parent1e18320
6 files changed+45−24
compare with main
unifiedsplit
internal/config/config.go+6−3 internal/llm/agent/agent.go−6 internal/llm/provider/anthropic.go+11−1 internal/llm/provider/gemini.go+21−5 internal/llm/provider/openai.go+7−2 internal/llm/provider/provider.go−7
internal/config/config.go+6−3
@@ -157,9 +157,12 @@ type Options struct {
157157}
158158
159159type PreferredModel struct {
160− ModelID string `json:"model_id"`
161− Provider provider.InferenceProvider `json:"provider"`
162− ReasoningEffort string `json:"reasoning_effort,omitempty"`
160+ ModelID string `json:"model_id"`
161+ Provider provider.InferenceProvider `json:"provider"`
162+ // Overrides the default reasoning effort for this model
163+ ReasoningEffort string `json:"reasoning_effort,omitempty"`
164+ // Overrides the default max tokens for this model
165+ MaxTokens int64 `json:"max_tokens,omitempty"`
163166}
164167
165168type PreferredModels struct {
internal/llm/agent/agent.go−6
@@ -147,7 +147,6 @@ func NewAgent(
147147 opts := []provider.ProviderClientOption{
148148 provider.WithModel(agentCfg.Model),
149149 provider.WithSystemMessage(prompt.GetPrompt(promptID, providerCfg.ID)),
150− provider.WithMaxTokens(model.DefaultMaxTokens),
151150 }
152151 agentProvider, err := provider.NewProviderV2(providerCfg, opts...)
153152 if err != nil {
@@ -184,7 +183,6 @@ func NewAgent(
184183 titleOpts := []provider.ProviderClientOption{
185184 provider.WithModel(config.SmallModel),
186185 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),
187− provider.WithMaxTokens(40),
188186 }
189187 titleProvider, err := provider.NewProviderV2(smallModelProviderCfg, titleOpts...)
190188 if err != nil {
@@ -193,7 +191,6 @@ func NewAgent(
193191 summarizeOpts := []provider.ProviderClientOption{
194192 provider.WithModel(config.SmallModel),
195193 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),
196− provider.WithMaxTokens(smallModel.DefaultMaxTokens),
197194 }
198195 summarizeProvider, err := provider.NewProviderV2(smallModelProviderCfg, summarizeOpts...)
199196 if err != nil {
@@ -832,7 +829,6 @@ func (a *agent) UpdateModel() error {
832829 opts := []provider.ProviderClientOption{
833830 provider.WithModel(a.agentCfg.Model),
834831 provider.WithSystemMessage(prompt.GetPrompt(promptID, currentProviderCfg.ID)),
835− provider.WithMaxTokens(model.DefaultMaxTokens),
836832 }
837833
838834 newProvider, err := provider.NewProviderV2(currentProviderCfg, opts...)
@@ -877,7 +873,6 @@ func (a *agent) UpdateModel() error {
877873 titleOpts := []provider.ProviderClientOption{
878874 provider.WithModel(config.SmallModel),
879875 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),
880− provider.WithMaxTokens(40),
881876 }
882877 newTitleProvider, err := provider.NewProviderV2(smallModelProviderCfg, titleOpts...)
883878 if err != nil {
@@ -888,7 +883,6 @@ func (a *agent) UpdateModel() error {
888883 summarizeOpts := []provider.ProviderClientOption{
889884 provider.WithModel(config.SmallModel),
890885 provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),
891− provider.WithMaxTokens(smallModel.DefaultMaxTokens),
892886 }
893887 newSummarizeProvider, err := provider.NewProviderV2(smallModelProviderCfg, summarizeOpts...)
894888 if err != nil {
internal/llm/provider/anthropic.go+11−1
@@ -164,9 +164,19 @@ func (a *anthropicClient) preparedMessages(messages []anthropic.MessageParam, to
164164 // }
165165 // }
166166
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+
167177 return anthropic.MessageNewParams{
168178 Model: anthropic.Model(model.ID),
169− MaxTokens: a.providerOptions.maxTokens,
179+ MaxTokens: maxTokens,
170180 Temperature: temperature,
171181 Messages: messages,
172182 Tools: tools,
internal/llm/provider/gemini.go+21−5
@@ -155,17 +155,26 @@ func (g *geminiClient) finishReason(reason genai.FinishReason) message.FinishRea
155155func (g *geminiClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
156156 // Convert messages
157157 geminiMessages := g.convertMessages(messages)
158−
158+ model := g.providerOptions.model(g.providerOptions.modelType)
159159 cfg := config.Get()
160160 if cfg.Options.Debug {
161161 jsonData, _ := json.Marshal(geminiMessages)
162162 logging.Debug("Prepared messages", "messages", string(jsonData))
163163 }
164164
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+ }
165174 history := geminiMessages[:len(geminiMessages)-1] // All but last message
166175 lastMsg := geminiMessages[len(geminiMessages)-1]
167176 config := &genai.GenerateContentConfig{
168− MaxOutputTokens: int32(g.providerOptions.maxTokens),
177+ MaxOutputTokens: int32(maxTokens),
169178 SystemInstruction: &genai.Content{
170179 Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},
171180 },
@@ -173,7 +182,6 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too
173182 if len(tools) > 0 {
174183 config.Tools = g.convertTools(tools)
175184 }
176− model := g.providerOptions.model(g.providerOptions.modelType)
177185 chat, _ := g.client.Chats.Create(ctx, model.ID, config, history)
178186
179187 attempts := 0
@@ -245,16 +253,25 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
245253 // Convert messages
246254 geminiMessages := g.convertMessages(messages)
247255
256+ model := g.providerOptions.model(g.providerOptions.modelType)
248257 cfg := config.Get()
249258 if cfg.Options.Debug {
250259 jsonData, _ := json.Marshal(geminiMessages)
251260 logging.Debug("Prepared messages", "messages", string(jsonData))
252261 }
253262
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+ }
254271 history := geminiMessages[:len(geminiMessages)-1] // All but last message
255272 lastMsg := geminiMessages[len(geminiMessages)-1]
256273 config := &genai.GenerateContentConfig{
257− MaxOutputTokens: int32(g.providerOptions.maxTokens),
274+ MaxOutputTokens: int32(maxTokens),
258275 SystemInstruction: &genai.Content{
259276 Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},
260277 },
@@ -262,7 +279,6 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
262279 if len(tools) > 0 {
263280 config.Tools = g.convertTools(tools)
264281 }
265− model := g.providerOptions.model(g.providerOptions.modelType)
266282 chat, _ := g.client.Chats.Create(ctx, model.ID, config, history)
267283
268284 attempts := 0
internal/llm/provider/openai.go+7−2
@@ -160,8 +160,13 @@ func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessagePar
160160 Messages: messages,
161161 Tools: tools,
162162 }
163+
164+ maxTokens := model.DefaultMaxTokens
165+ if modelConfig.MaxTokens > 0 {
166+ maxTokens = modelConfig.MaxTokens
167+ }
163168 if model.CanReason {
164− params.MaxCompletionTokens = openai.Int(o.providerOptions.maxTokens)
169+ params.MaxCompletionTokens = openai.Int(maxTokens)
165170 switch reasoningEffort {
166171 case "low":
167172 params.ReasoningEffort = shared.ReasoningEffortLow
@@ -173,7 +178,7 @@ func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessagePar
173178 params.ReasoningEffort = shared.ReasoningEffortMedium
174179 }
175180 } else {
176− params.MaxTokens = openai.Int(o.providerOptions.maxTokens)
181+ params.MaxTokens = openai.Int(maxTokens)
177182 }
178183
179184 return params
internal/llm/provider/provider.go−7
@@ -64,7 +64,6 @@ type providerClientOptions struct {
6464 modelType config.ModelType
6565 model func(config.ModelType) config.Model
6666 disableCache bool
67− maxTokens int64
6867 systemMessage string
6968 extraHeaders map[string]string
7069 extraParams map[string]string
@@ -121,12 +120,6 @@ func WithDisableCache(disableCache bool) ProviderClientOption {
121120 }
122121}
123122
124−func WithMaxTokens(maxTokens int64) ProviderClientOption {
125− return func(options *providerClientOptions) {
126− options.maxTokens = maxTokens
127− }
128−}
129−
130123func WithSystemMessage(systemMessage string) ProviderClientOption {
131124 return func(options *providerClientOptions) {
132125 options.systemMessage = systemMessage