8c06123 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

chore: change how we handle max tokens for anthropic

8c06123Kujtim Hoxha1 year ago

internal/config/config.go−38

compare with main

unifiedsplit

@@ -890,44 +890,6 @@ func GetAgentModel(agentID AgentID) Model {

890 return Model{}

891}

892

893−// GetAgentEffectiveMaxTokens returns the effective max tokens for an agent,

894−// considering any overrides from the preferred model configuration

895−func GetAgentEffectiveMaxTokens(agentID AgentID) int64 {

896− cfg := Get()

897− agent, ok := cfg.Agents[agentID]

898− if !ok {

899− logging.Error("Agent not found", "agent_id", agentID)

900− return 0

901− }

902−

903− var preferredModel PreferredModel

904− switch agent.Model {

905− case LargeModel:

906− preferredModel = cfg.Models.Large

907− case SmallModel:

908− preferredModel = cfg.Models.Small

909− default:

910− logging.Warn("Unknown model type for agent", "agent_id", agentID, "model_type", agent.Model)

911− preferredModel = cfg.Models.Large // Fallback to large model

912− }

913−

914− // Get the base model configuration

915− baseModel := GetAgentModel(agentID)

916− if baseModel.ID == "" {

917− return 0

918− }

919−

920− // Start with the default max tokens from the base model

921− maxTokens := baseModel.DefaultMaxTokens

922−

923− // Override with preferred model max tokens if set

924− if preferredModel.MaxTokens > 0 {

925− maxTokens = preferredModel.MaxTokens

926− }

927−

928− return maxTokens

929−}

930−

931func GetAgentProvider(agentID AgentID) ProviderConfig {

893func GetAgentProvider(agentID AgentID) ProviderConfig {

932 cfg := Get()

894 cfg := Get()

933 agent, ok := cfg.Agents[agentID]

895 agent, ok := cfg.Agents[agentID]

internal/llm/agent/agent.go−5

@@ -50,7 +50,6 @@ type AgentEvent struct {

50type Service interface {

51 pubsub.Suscriber[AgentEvent]

52 Model() config.Model

53− EffectiveMaxTokens() int64

54 Run(ctx context.Context, sessionID string, content string, attachments ...message.Attachment) (<-chan AgentEvent, error)

53 Run(ctx context.Context, sessionID string, content string, attachments ...message.Attachment) (<-chan AgentEvent, error)

55 Cancel(sessionID string)

54 Cancel(sessionID string)

56 CancelAll()

55 CancelAll()

@@ -230,10 +229,6 @@ func (a *agent) Model() config.Model {

230 return config.GetAgentModel(a.agentCfg.ID)

229 return config.GetAgentModel(a.agentCfg.ID)

231}

230}

232

231

233−func (a *agent) EffectiveMaxTokens() int64 {

234− return config.GetAgentEffectiveMaxTokens(a.agentCfg.ID)

235−}

236−

237func (a *agent) Cancel(sessionID string) {

232func (a *agent) Cancel(sessionID string) {

238 // Cancel regular requests

233 // Cancel regular requests

239 if cancelFunc, exists := a.activeRequests.LoadAndDelete(sessionID); exists {

234 if cancelFunc, exists := a.activeRequests.LoadAndDelete(sessionID); exists {

internal/llm/provider/anthropic.go+61−13

@@ -6,6 +6,8 @@ import (

6 "errors"

7 "fmt"

8 "io"

9+ "regexp"

10+ "strconv"

9 "time"

11 "time"

11 "github.com/anthropics/anthropic-sdk-go"

13 "github.com/anthropics/anthropic-sdk-go"

@@ -19,9 +21,10 @@ import (

19)

21)

21type anthropicClient struct {

23type anthropicClient struct {

22− providerOptions providerClientOptions

24+ providerOptions providerClientOptions

23− useBedrock bool

25+ useBedrock bool

24− client anthropic.Client

26+ client anthropic.Client

27+ adjustedMaxTokens int // Used when context limit is hit

25}

28}

27type AnthropicClient ProviderClient

30type AnthropicClient ProviderClient

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

171 maxTokens = a.providerOptions.maxTokens

174 maxTokens = a.providerOptions.maxTokens

172 }

175 }

173

176

177+ // Use adjusted max tokens if context limit was hit

178+ if a.adjustedMaxTokens > 0 {

179+ maxTokens = int64(a.adjustedMaxTokens)

180+ }

181+

174 return anthropic.MessageNewParams{

182 return anthropic.MessageNewParams{

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

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

176 MaxTokens: maxTokens,

184 MaxTokens: maxTokens,

@@ -190,16 +198,18 @@ func (a *anthropicClient) preparedMessages(messages []anthropic.MessageParam, to

190}

198}

191

199

192func (a *anthropicClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (response *ProviderResponse, err error) {

200func (a *anthropicClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (response *ProviderResponse, err error) {

193− preparedMessages := a.preparedMessages(a.convertMessages(messages), a.convertTools(tools))

194 cfg := config.Get()

201 cfg := config.Get()

195− if cfg.Options.Debug {

196− jsonData, _ := json.Marshal(preparedMessages)

197− logging.Debug("Prepared messages", "messages", string(jsonData))

198− }

199

202

200 attempts := 0

203 attempts := 0

201 for {

204 for {

202 attempts++

205 attempts++

206+ // Prepare messages on each attempt in case max_tokens was adjusted

207+ preparedMessages := a.preparedMessages(a.convertMessages(messages), a.convertTools(tools))

208+ if cfg.Options.Debug {

209+ jsonData, _ := json.Marshal(preparedMessages)

210+ logging.Debug("Prepared messages", "messages", string(jsonData))

211+ }

212+

203 anthropicResponse, err := a.client.Messages.New(

213 anthropicResponse, err := a.client.Messages.New(

204 ctx,

214 ctx,

205 preparedMessages,

215 preparedMessages,

@@ -239,17 +249,19 @@ func (a *anthropicClient) send(ctx context.Context, messages []message.Message,

239}

249}

240

250

241func (a *anthropicClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {

251func (a *anthropicClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {

242− preparedMessages := a.preparedMessages(a.convertMessages(messages), a.convertTools(tools))

243 cfg := config.Get()

252 cfg := config.Get()

244− if cfg.Options.Debug {

245− // jsonData, _ := json.Marshal(preparedMessages)

246− // logging.Debug("Prepared messages", "messages", string(jsonData))

247− }

248 attempts := 0

253 attempts := 0

249 eventChan := make(chan ProviderEvent)

254 eventChan := make(chan ProviderEvent)

250 go func() {

255 go func() {

251 for {

256 for {

252 attempts++

257 attempts++

258+ // Prepare messages on each attempt in case max_tokens was adjusted

259+ preparedMessages := a.preparedMessages(a.convertMessages(messages), a.convertTools(tools))

260+ if cfg.Options.Debug {

261+ jsonData, _ := json.Marshal(preparedMessages)

262+ logging.Debug("Prepared messages", "messages", string(jsonData))

263+ }

264+

253 anthropicStream := a.client.Messages.NewStreaming(

265 anthropicStream := a.client.Messages.NewStreaming(

254 ctx,

266 ctx,

255 preparedMessages,

267 preparedMessages,

@@ -395,6 +407,15 @@ func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, err

395 return true, 0, nil

407 return true, 0, nil

396 }

408 }

397

409

410+ // Handle context limit exceeded error (400 Bad Request)

411+ if apiErr.StatusCode == 400 {

412+ if adjusted, ok := a.handleContextLimitError(apiErr); ok {

413+ a.adjustedMaxTokens = adjusted

414+ logging.Debug("Adjusted max_tokens due to context limit", "new_max_tokens", adjusted)

415+ return true, 0, nil

416+ }

417+ }

418+

398 if apiErr.StatusCode != 429 && apiErr.StatusCode != 529 {

419 if apiErr.StatusCode != 429 && apiErr.StatusCode != 529 {

399 return false, 0, err

420 return false, 0, err

400 }

421 }

@@ -413,6 +434,33 @@ func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, err

413 return true, int64(retryMs), nil

434 return true, int64(retryMs), nil

414}

435}

415

436

437+// handleContextLimitError parses context limit error and returns adjusted max_tokens

438+func (a *anthropicClient) handleContextLimitError(apiErr *anthropic.Error) (int, bool) {

439+ // Parse error message like: "input length and max_tokens exceed context limit: 154978 + 50000 > 200000"

440+ errorMsg := apiErr.Error()

441+ re := regexp.MustCompile(`input length and max_tokens exceed context limit: (\d+) \+ (\d+) > (\d+)`)

442+ matches := re.FindStringSubmatch(errorMsg)

443+

444+ if len(matches) != 4 {

445+ return 0, false

446+ }

447+

448+ inputTokens, err1 := strconv.Atoi(matches[1])

449+ contextLimit, err2 := strconv.Atoi(matches[3])

450+

451+ if err1 != nil || err2 != nil {

452+ return 0, false

453+ }

454+

455+ // Calculate safe max_tokens with a buffer of 1000 tokens

456+ safeMaxTokens := contextLimit - inputTokens - 1000

457+

458+ // Ensure we don't go below a minimum threshold

459+ safeMaxTokens = max(safeMaxTokens, 1000)

460+

461+ return safeMaxTokens, true

462+}

463+

416func (a *anthropicClient) toolCalls(msg anthropic.Message) []message.ToolCall {

464func (a *anthropicClient) toolCalls(msg anthropic.Message) []message.ToolCall {

417 var toolCalls []message.ToolCall

465 var toolCalls []message.ToolCall

418

466

internal/tui/tui.go+2−18

@@ -228,24 +228,8 @@ func (a *appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

228 if err == nil {

229 model := a.app.CoderAgent.Model()

230 contextWindow := model.ContextWindow

231− usedTokens := session.CompletionTokens + session.PromptTokens

231+ tokens := session.CompletionTokens + session.PromptTokens

232− remainingTokens := contextWindow - usedTokens

233−

234− // Get effective max tokens for this agent (considering overrides)

235− maxTokens := a.app.CoderAgent.EffectiveMaxTokens()

236−

237− // Apply 10% margin to max tokens

238− maxTokensWithMargin := int64(float64(maxTokens) * 1.1)

239−

240− // Trigger auto-summarize if remaining tokens < max tokens + 10% margin

241− // Also ensure we have a reasonable minimum threshold to avoid too-frequent summaries

242− minThreshold := int64(1000) // Minimum 1000 tokens remaining before triggering

243− if maxTokensWithMargin < minThreshold {

244− maxTokensWithMargin = minThreshold

245− }

246−

247− if remainingTokens < maxTokensWithMargin && !config.Get().Options.DisableAutoSummarize {

248− // Show compact confirmation dialog

232+ if (tokens >= int64(float64(contextWindow)*0.95)) && !config.Get().Options.DisableAutoSummarize { // Show compact confirmation dialog

249 cmds = append(cmds, util.CmdHandler(dialogs.OpenDialogMsg{

233 cmds = append(cmds, util.CmdHandler(dialogs.OpenDialogMsg{

250 Model: compact.NewCompactDialogCmp(a.app.CoderAgent, a.selectedSessionID, false),

234 Model: compact.NewCompactDialogCmp(a.app.CoderAgent, a.selectedSessionID, false),

251 }))

235 }))