03871a5 · crush · git.iko.soy

git.iko.soy/crush

chore: make it possible to override maxTokens

we still need to handle the case where the max tokens is more than 5% of the total context window, this can cause the endpoint to error.

commit03871a565bb8f51d274148c266fb4ae5a7a89c9a authorKujtim Hoxha <kujtimii.h@gmail.com> dateSat 28 Jun 2025, 12:06 (1 year ago) parenteb3550a

5 files changed+34−8

compare with main

unifiedsplit

internal/llm/agent/agent.go+11−7 internal/llm/provider/anthropic.go+5 internal/llm/provider/gemini.go+5 internal/llm/provider/openai.go+5 internal/llm/provider/provider.go+8−1

internal/llm/agent/agent.go+11−7

@@ -148,7 +148,7 @@ func NewAgent(

148148 provider.WithModel(agentCfg.Model),

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

150150 }

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

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

152152 if err != nil {

153153 return nil, err

154154 }

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

184184 provider.WithModel(config.SmallModel),

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

186186 }

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

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

188188 if err != nil {

189189 return nil, err

190190 }

@@ -192,7 +192,7 @@ func NewAgent(

192192 provider.WithModel(config.SmallModel),

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

194194 }

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

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

196196 if err != nil {

197197 return nil, err

198198 }

@@ -277,7 +277,9 @@ func (a *agent) generateTitle(ctx context.Context, sessionID string, content str

277277 if err != nil {

278278 return err

279279 }

280− parts := []message.ContentPart{message.TextContent{Text: content}}

280+ parts := []message.ContentPart{message.TextContent{

281+ Text: fmt.Sprintf("Generate a concise title for the following content:\n\n%s", content),

282+ }}

281283

282284 // Use streaming approach like summarization

283285 response := a.titleProvider.StreamResponse(

@@ -831,7 +833,7 @@ func (a *agent) UpdateModel() error {

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

832834 }

833835

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

836+ newProvider, err := provider.NewProvider(currentProviderCfg, opts...)

835837 if err != nil {

836838 return fmt.Errorf("failed to create new provider: %w", err)

837839 }

@@ -873,8 +875,10 @@ func (a *agent) UpdateModel() error {

873875 titleOpts := []provider.ProviderClientOption{

874876 provider.WithModel(config.SmallModel),

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

878+ // We want the title to be short, so we limit the max tokens

879+ provider.WithMaxTokens(40),

876880 }

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

881+ newTitleProvider, err := provider.NewProvider(smallModelProviderCfg, titleOpts...)

878882 if err != nil {

879883 return fmt.Errorf("failed to create new title provider: %w", err)

880884 }

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

884888 provider.WithModel(config.SmallModel),

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

886890 }

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

891+ newSummarizeProvider, err := provider.NewProvider(smallModelProviderCfg, summarizeOpts...)

888892 if err != nil {

889893 return fmt.Errorf("failed to create new summarize provider: %w", err)

890894 }

internal/llm/provider/anthropic.go+5

@@ -174,6 +174,11 @@ func (a *anthropicClient) preparedMessages(messages []anthropic.MessageParam, to

174174 maxTokens = modelConfig.MaxTokens

175175 }

176176

177+ // Override max tokens if set in provider options

178+ if a.providerOptions.maxTokens > 0 {

179+ maxTokens = a.providerOptions.maxTokens

180+ }

181+

177182 return anthropic.MessageNewParams{

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

179184 MaxTokens: maxTokens,

internal/llm/provider/gemini.go+5

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

268268 if modelConfig.MaxTokens > 0 {

269269 maxTokens = modelConfig.MaxTokens

270270 }

271+

272+ // Override max tokens if set in provider options

273+ if g.providerOptions.maxTokens > 0 {

274+ maxTokens = g.providerOptions.maxTokens

275+ }

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

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

273278 config := &genai.GenerateContentConfig{

internal/llm/provider/openai.go+5

@@ -165,6 +165,11 @@ func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessagePar

165165 if modelConfig.MaxTokens > 0 {

166166 maxTokens = modelConfig.MaxTokens

167167 }

168+

169+ // Override max tokens if set in provider options

170+ if o.providerOptions.maxTokens > 0 {

171+ maxTokens = o.providerOptions.maxTokens

172+ }

168173 if model.CanReason {

169174 params.MaxCompletionTokens = openai.Int(maxTokens)

170175 switch reasoningEffort {

internal/llm/provider/provider.go+8−1

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

6565 model func(config.ModelType) config.Model

6666 disableCache bool

6767 systemMessage string

68+ maxTokens int64

6869 extraHeaders map[string]string

6970 extraParams map[string]string

7071}

@@ -126,7 +127,13 @@ func WithSystemMessage(systemMessage string) ProviderClientOption {

126127 }

127128}

128129

129−func NewProviderV2(cfg config.ProviderConfig, opts ...ProviderClientOption) (Provider, error) {

130+func WithMaxTokens(maxTokens int64) ProviderClientOption {

131+ return func(options *providerClientOptions) {

132+ options.maxTokens = maxTokens

133+ }

134+}

135+

136+func NewProvider(cfg config.ProviderConfig, opts ...ProviderClientOption) (Provider, error) {

130137 clientOptions := providerClientOptions{

131138 baseURL: cfg.BaseURL,

132139 apiKey: cfg.APIKey,