03871a5 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

chore: make it possible to override maxTokens

03871a5Kujtim Hoxha1 year ago

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

compare with main

unifiedsplit

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

148 provider.WithModel(agentCfg.Model),

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

150 }

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

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

152 if err != nil {

153 return nil, err

154 }

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

184 provider.WithModel(config.SmallModel),

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

186 }

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

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

188 if err != nil {

189 return nil, err

190 }

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

192 provider.WithModel(config.SmallModel),

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

194 }

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

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

196 if err != nil {

197 return nil, err

198 }

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

277 if err != nil {

278 return err

279 }

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+ }}

281

283

282 // Use streaming approach like summarization

284 // Use streaming approach like summarization

283 response := a.titleProvider.StreamResponse(

285 response := a.titleProvider.StreamResponse(

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

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

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

832 }

834 }

833

835

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

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

835 if err != nil {

837 if err != nil {

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

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

837 }

839 }

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

873 titleOpts := []provider.ProviderClientOption{

875 titleOpts := []provider.ProviderClientOption{

874 provider.WithModel(config.SmallModel),

876 provider.WithModel(config.SmallModel),

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

877 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),

876 }

880 }

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

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

878 if err != nil {

882 if err != nil {

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

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

880 }

884 }

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

884 provider.WithModel(config.SmallModel),

888 provider.WithModel(config.SmallModel),

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

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

886 }

890 }

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

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

888 if err != nil {

892 if err != nil {

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

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

890 }

894 }

internal/llm/provider/anthropic.go+5

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

174 maxTokens = modelConfig.MaxTokens

175 }

176

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

178+ if a.providerOptions.maxTokens > 0 {

179+ maxTokens = a.providerOptions.maxTokens

180+ }

181+

177 return anthropic.MessageNewParams{

182 return anthropic.MessageNewParams{

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

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

179 MaxTokens: maxTokens,

184 MaxTokens: maxTokens,

internal/llm/provider/gemini.go+5

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

268 if modelConfig.MaxTokens > 0 {

269 maxTokens = modelConfig.MaxTokens

270 }

271+

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

273+ if g.providerOptions.maxTokens > 0 {

274+ maxTokens = g.providerOptions.maxTokens

275+ }

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

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

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

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

273 config := &genai.GenerateContentConfig{

278 config := &genai.GenerateContentConfig{

internal/llm/provider/openai.go+5

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

165 if modelConfig.MaxTokens > 0 {

166 maxTokens = modelConfig.MaxTokens

167 }

168+

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

170+ if o.providerOptions.maxTokens > 0 {

171+ maxTokens = o.providerOptions.maxTokens

172+ }

168 if model.CanReason {

173 if model.CanReason {

169 params.MaxCompletionTokens = openai.Int(maxTokens)

174 params.MaxCompletionTokens = openai.Int(maxTokens)

170 switch reasoningEffort {

175 switch reasoningEffort {

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

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

65 model func(config.ModelType) config.Model

66 disableCache bool

67 systemMessage string

68+ maxTokens int64

68 extraHeaders map[string]string

69 extraHeaders map[string]string

69 extraParams map[string]string

70 extraParams map[string]string

70}

71}

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

126 }

127 }

127}

128}

128

129

130+func WithMaxTokens(maxTokens int64) ProviderClientOption {

131+ return func(options *providerClientOptions) {

132+ options.maxTokens = maxTokens

133+ }

134+}

135+

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

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

130 clientOptions := providerClientOptions{

137 clientOptions := providerClientOptions{

131 baseURL: cfg.BaseURL,

138 baseURL: cfg.BaseURL,

132 apiKey: cfg.APIKey,

139 apiKey: cfg.APIKey,