1package provider
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"errors"
  7	"fmt"
  8	"io"
  9	"strings"
 10	"time"
 11
 12	"github.com/anthropics/anthropic-sdk-go"
 13	"github.com/anthropics/anthropic-sdk-go/bedrock"
 14	"github.com/anthropics/anthropic-sdk-go/option"
 15	"github.com/charmbracelet/crush/internal/config"
 16	"github.com/charmbracelet/crush/internal/llm/models"
 17	"github.com/charmbracelet/crush/internal/llm/tools"
 18	"github.com/charmbracelet/crush/internal/logging"
 19	"github.com/charmbracelet/crush/internal/message"
 20)
 21
 22type anthropicOptions struct {
 23	useBedrock   bool
 24	disableCache bool
 25	shouldThink  func(userMessage string) bool
 26}
 27
 28type AnthropicOption func(*anthropicOptions)
 29
 30type anthropicClient struct {
 31	providerOptions providerClientOptions
 32	options         anthropicOptions
 33	client          anthropic.Client
 34}
 35
 36type AnthropicClient ProviderClient
 37
 38func newAnthropicClient(opts providerClientOptions) AnthropicClient {
 39	anthropicOpts := anthropicOptions{}
 40	for _, o := range opts.anthropicOptions {
 41		o(&anthropicOpts)
 42	}
 43
 44	anthropicClientOptions := []option.RequestOption{}
 45	if opts.apiKey != "" {
 46		anthropicClientOptions = append(anthropicClientOptions, option.WithAPIKey(opts.apiKey))
 47	}
 48	if anthropicOpts.useBedrock {
 49		anthropicClientOptions = append(anthropicClientOptions, bedrock.WithLoadDefaultConfig(context.Background()))
 50	}
 51
 52	client := anthropic.NewClient(anthropicClientOptions...)
 53	return &anthropicClient{
 54		providerOptions: opts,
 55		options:         anthropicOpts,
 56		client:          client,
 57	}
 58}
 59
 60func (a *anthropicClient) convertMessages(messages []message.Message) (anthropicMessages []anthropic.MessageParam) {
 61	for i, msg := range messages {
 62		cache := false
 63		if i > len(messages)-3 {
 64			cache = true
 65		}
 66		switch msg.Role {
 67		case message.User:
 68			content := anthropic.NewTextBlock(msg.Content().String())
 69			if cache && !a.options.disableCache {
 70				content.OfText.CacheControl = anthropic.CacheControlEphemeralParam{
 71					Type: "ephemeral",
 72				}
 73			}
 74			var contentBlocks []anthropic.ContentBlockParamUnion
 75			contentBlocks = append(contentBlocks, content)
 76			for _, binaryContent := range msg.BinaryContent() {
 77				base64Image := binaryContent.String(models.ProviderAnthropic)
 78				imageBlock := anthropic.NewImageBlockBase64(binaryContent.MIMEType, base64Image)
 79				contentBlocks = append(contentBlocks, imageBlock)
 80			}
 81			anthropicMessages = append(anthropicMessages, anthropic.NewUserMessage(contentBlocks...))
 82
 83		case message.Assistant:
 84			blocks := []anthropic.ContentBlockParamUnion{}
 85			if msg.Content().String() != "" {
 86				content := anthropic.NewTextBlock(msg.Content().String())
 87				if cache && !a.options.disableCache {
 88					content.OfText.CacheControl = anthropic.CacheControlEphemeralParam{
 89						Type: "ephemeral",
 90					}
 91				}
 92				blocks = append(blocks, content)
 93			}
 94
 95			for _, toolCall := range msg.ToolCalls() {
 96				var inputMap map[string]any
 97				err := json.Unmarshal([]byte(toolCall.Input), &inputMap)
 98				if err != nil {
 99					continue
100				}
101				blocks = append(blocks, anthropic.NewToolUseBlock(toolCall.ID, inputMap, toolCall.Name))
102			}
103
104			if len(blocks) == 0 {
105				logging.Warn("There is a message without content, investigate, this should not happen")
106				continue
107			}
108			anthropicMessages = append(anthropicMessages, anthropic.NewAssistantMessage(blocks...))
109
110		case message.Tool:
111			results := make([]anthropic.ContentBlockParamUnion, len(msg.ToolResults()))
112			for i, toolResult := range msg.ToolResults() {
113				results[i] = anthropic.NewToolResultBlock(toolResult.ToolCallID, toolResult.Content, toolResult.IsError)
114			}
115			anthropicMessages = append(anthropicMessages, anthropic.NewUserMessage(results...))
116		}
117	}
118	return
119}
120
121func (a *anthropicClient) convertTools(tools []tools.BaseTool) []anthropic.ToolUnionParam {
122	anthropicTools := make([]anthropic.ToolUnionParam, len(tools))
123
124	for i, tool := range tools {
125		info := tool.Info()
126		toolParam := anthropic.ToolParam{
127			Name:        info.Name,
128			Description: anthropic.String(info.Description),
129			InputSchema: anthropic.ToolInputSchemaParam{
130				Properties: info.Parameters,
131				// TODO: figure out how we can tell claude the required fields?
132			},
133		}
134
135		if i == len(tools)-1 && !a.options.disableCache {
136			toolParam.CacheControl = anthropic.CacheControlEphemeralParam{
137				Type: "ephemeral",
138			}
139		}
140
141		anthropicTools[i] = anthropic.ToolUnionParam{OfTool: &toolParam}
142	}
143
144	return anthropicTools
145}
146
147func (a *anthropicClient) finishReason(reason string) message.FinishReason {
148	switch reason {
149	case "end_turn":
150		return message.FinishReasonEndTurn
151	case "max_tokens":
152		return message.FinishReasonMaxTokens
153	case "tool_use":
154		return message.FinishReasonToolUse
155	case "stop_sequence":
156		return message.FinishReasonEndTurn
157	default:
158		return message.FinishReasonUnknown
159	}
160}
161
162func (a *anthropicClient) preparedMessages(messages []anthropic.MessageParam, tools []anthropic.ToolUnionParam) anthropic.MessageNewParams {
163	var thinkingParam anthropic.ThinkingConfigParamUnion
164	lastMessage := messages[len(messages)-1]
165	isUser := lastMessage.Role == anthropic.MessageParamRoleUser
166	messageContent := ""
167	temperature := anthropic.Float(0)
168	if isUser {
169		for _, m := range lastMessage.Content {
170			if m.OfText != nil && m.OfText.Text != "" {
171				messageContent = m.OfText.Text
172			}
173		}
174		if messageContent != "" && a.options.shouldThink != nil && a.options.shouldThink(messageContent) {
175			thinkingParam = anthropic.ThinkingConfigParamOfEnabled(int64(float64(a.providerOptions.maxTokens) * 0.8))
176			temperature = anthropic.Float(1)
177		}
178	}
179
180	return anthropic.MessageNewParams{
181		Model:       anthropic.Model(a.providerOptions.model.APIModel),
182		MaxTokens:   a.providerOptions.maxTokens,
183		Temperature: temperature,
184		Messages:    messages,
185		Tools:       tools,
186		Thinking:    thinkingParam,
187		System: []anthropic.TextBlockParam{
188			{
189				Text: a.providerOptions.systemMessage,
190				CacheControl: anthropic.CacheControlEphemeralParam{
191					Type: "ephemeral",
192				},
193			},
194		},
195	}
196}
197
198func (a *anthropicClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (response *ProviderResponse, err error) {
199	preparedMessages := a.preparedMessages(a.convertMessages(messages), a.convertTools(tools))
200	cfg := config.Get()
201	if cfg.Debug {
202		jsonData, _ := json.Marshal(preparedMessages)
203		logging.Debug("Prepared messages", "messages", string(jsonData))
204	}
205
206	attempts := 0
207	for {
208		attempts++
209		anthropicResponse, err := a.client.Messages.New(
210			ctx,
211			preparedMessages,
212		)
213		// If there is an error we are going to see if we can retry the call
214		if err != nil {
215			logging.Error("Error in Anthropic API call", "error", err)
216			retry, after, retryErr := a.shouldRetry(attempts, err)
217			if retryErr != nil {
218				return nil, retryErr
219			}
220			if retry {
221				logging.WarnPersist(fmt.Sprintf("Retrying due to rate limit... attempt %d of %d", attempts, maxRetries), logging.PersistTimeArg, time.Millisecond*time.Duration(after+100))
222				select {
223				case <-ctx.Done():
224					return nil, ctx.Err()
225				case <-time.After(time.Duration(after) * time.Millisecond):
226					continue
227				}
228			}
229			return nil, retryErr
230		}
231
232		content := ""
233		for _, block := range anthropicResponse.Content {
234			if text, ok := block.AsAny().(anthropic.TextBlock); ok {
235				content += text.Text
236			}
237		}
238
239		return &ProviderResponse{
240			Content:   content,
241			ToolCalls: a.toolCalls(*anthropicResponse),
242			Usage:     a.usage(*anthropicResponse),
243		}, nil
244	}
245}
246
247func (a *anthropicClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {
248	preparedMessages := a.preparedMessages(a.convertMessages(messages), a.convertTools(tools))
249	cfg := config.Get()
250	if cfg.Debug {
251		// jsonData, _ := json.Marshal(preparedMessages)
252		// logging.Debug("Prepared messages", "messages", string(jsonData))
253	}
254	attempts := 0
255	eventChan := make(chan ProviderEvent)
256	go func() {
257		for {
258			attempts++
259			anthropicStream := a.client.Messages.NewStreaming(
260				ctx,
261				preparedMessages,
262			)
263			accumulatedMessage := anthropic.Message{}
264
265			currentToolCallID := ""
266			for anthropicStream.Next() {
267				event := anthropicStream.Current()
268				err := accumulatedMessage.Accumulate(event)
269				if err != nil {
270					logging.Warn("Error accumulating message", "error", err)
271					continue
272				}
273
274				switch event := event.AsAny().(type) {
275				case anthropic.ContentBlockStartEvent:
276					switch event.ContentBlock.Type {
277					case "text":
278						eventChan <- ProviderEvent{Type: EventContentStart}
279					case "tool_use":
280						currentToolCallID = event.ContentBlock.ID
281						eventChan <- ProviderEvent{
282							Type: EventToolUseStart,
283							ToolCall: &message.ToolCall{
284								ID:       event.ContentBlock.ID,
285								Name:     event.ContentBlock.Name,
286								Finished: false,
287							},
288						}
289					}
290
291				case anthropic.ContentBlockDeltaEvent:
292					if event.Delta.Type == "thinking_delta" && event.Delta.Thinking != "" {
293						eventChan <- ProviderEvent{
294							Type:     EventThinkingDelta,
295							Thinking: event.Delta.Thinking,
296						}
297					} else if event.Delta.Type == "text_delta" && event.Delta.Text != "" {
298						eventChan <- ProviderEvent{
299							Type:    EventContentDelta,
300							Content: event.Delta.Text,
301						}
302					} else if event.Delta.Type == "input_json_delta" {
303						if currentToolCallID != "" {
304							eventChan <- ProviderEvent{
305								Type: EventToolUseDelta,
306								ToolCall: &message.ToolCall{
307									ID:       currentToolCallID,
308									Finished: false,
309									Input:    event.Delta.PartialJSON,
310								},
311							}
312						}
313					}
314				case anthropic.ContentBlockStopEvent:
315					if currentToolCallID != "" {
316						eventChan <- ProviderEvent{
317							Type: EventToolUseStop,
318							ToolCall: &message.ToolCall{
319								ID: currentToolCallID,
320							},
321						}
322						currentToolCallID = ""
323					} else {
324						eventChan <- ProviderEvent{Type: EventContentStop}
325					}
326
327				case anthropic.MessageStopEvent:
328					content := ""
329					for _, block := range accumulatedMessage.Content {
330						if text, ok := block.AsAny().(anthropic.TextBlock); ok {
331							content += text.Text
332						}
333					}
334
335					eventChan <- ProviderEvent{
336						Type: EventComplete,
337						Response: &ProviderResponse{
338							Content:      content,
339							ToolCalls:    a.toolCalls(accumulatedMessage),
340							Usage:        a.usage(accumulatedMessage),
341							FinishReason: a.finishReason(string(accumulatedMessage.StopReason)),
342						},
343						Content: content,
344					}
345				}
346			}
347
348			err := anthropicStream.Err()
349			if err == nil || errors.Is(err, io.EOF) {
350				close(eventChan)
351				return
352			}
353			// If there is an error we are going to see if we can retry the call
354			retry, after, retryErr := a.shouldRetry(attempts, err)
355			if retryErr != nil {
356				eventChan <- ProviderEvent{Type: EventError, Error: retryErr}
357				close(eventChan)
358				return
359			}
360			if retry {
361				logging.WarnPersist(fmt.Sprintf("Retrying due to rate limit... attempt %d of %d", attempts, maxRetries), logging.PersistTimeArg, time.Millisecond*time.Duration(after+100))
362				select {
363				case <-ctx.Done():
364					// context cancelled
365					if ctx.Err() != nil {
366						eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
367					}
368					close(eventChan)
369					return
370				case <-time.After(time.Duration(after) * time.Millisecond):
371					continue
372				}
373			}
374			if ctx.Err() != nil {
375				eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
376			}
377
378			close(eventChan)
379			return
380		}
381	}()
382	return eventChan
383}
384
385func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, error) {
386	var apierr *anthropic.Error
387	if !errors.As(err, &apierr) {
388		return false, 0, err
389	}
390
391	if apierr.StatusCode != 429 && apierr.StatusCode != 529 {
392		return false, 0, err
393	}
394
395	if attempts > maxRetries {
396		return false, 0, fmt.Errorf("maximum retry attempts reached for rate limit: %d retries", maxRetries)
397	}
398
399	retryMs := 0
400	retryAfterValues := apierr.Response.Header.Values("Retry-After")
401
402	backoffMs := 2000 * (1 << (attempts - 1))
403	jitterMs := int(float64(backoffMs) * 0.2)
404	retryMs = backoffMs + jitterMs
405	if len(retryAfterValues) > 0 {
406		if _, err := fmt.Sscanf(retryAfterValues[0], "%d", &retryMs); err == nil {
407			retryMs = retryMs * 1000
408		}
409	}
410	return true, int64(retryMs), nil
411}
412
413func (a *anthropicClient) toolCalls(msg anthropic.Message) []message.ToolCall {
414	var toolCalls []message.ToolCall
415
416	for _, block := range msg.Content {
417		switch variant := block.AsAny().(type) {
418		case anthropic.ToolUseBlock:
419			toolCall := message.ToolCall{
420				ID:       variant.ID,
421				Name:     variant.Name,
422				Input:    string(variant.Input),
423				Type:     string(variant.Type),
424				Finished: true,
425			}
426			toolCalls = append(toolCalls, toolCall)
427		}
428	}
429
430	return toolCalls
431}
432
433func (a *anthropicClient) usage(msg anthropic.Message) TokenUsage {
434	return TokenUsage{
435		InputTokens:         msg.Usage.InputTokens,
436		OutputTokens:        msg.Usage.OutputTokens,
437		CacheCreationTokens: msg.Usage.CacheCreationInputTokens,
438		CacheReadTokens:     msg.Usage.CacheReadInputTokens,
439	}
440}
441
442func WithAnthropicBedrock(useBedrock bool) AnthropicOption {
443	return func(options *anthropicOptions) {
444		options.useBedrock = useBedrock
445	}
446}
447
448func WithAnthropicDisableCache() AnthropicOption {
449	return func(options *anthropicOptions) {
450		options.disableCache = true
451	}
452}
453
454func DefaultShouldThinkFn(s string) bool {
455	return strings.Contains(strings.ToLower(s), "think")
456}
457
458func WithAnthropicShouldThinkFn(fn func(string) bool) AnthropicOption {
459	return func(options *anthropicOptions) {
460		options.shouldThink = fn
461	}
462}