Files
agentic-coding-workflow/internal/infra/llm/anthropic.go
T

155 lines
4.9 KiB
Go

package llm
import (
"context"
"encoding/base64"
"fmt"
"github.com/anthropics/anthropic-sdk-go"
"github.com/anthropics/anthropic-sdk-go/option"
)
// Compile-time assertion: AnthropicClient must satisfy LLMClient.
var _ LLMClient = (*AnthropicClient)(nil)
// AnthropicClient 通过官方 SDK 调用 Anthropic Messages API。
// baseURL 可改写以走 LiteLLM/OpenRouter 等代理。
type AnthropicClient struct {
sdk anthropic.Client
}
// NewAnthropic 构造客户端。baseURL 可空(用 SDK 默认)。
func NewAnthropic(apiKey, baseURL string) (*AnthropicClient, error) {
opts := []option.RequestOption{option.WithAPIKey(apiKey)}
if baseURL != "" {
opts = append(opts, option.WithBaseURL(baseURL))
}
c := anthropic.NewClient(opts...)
return &AnthropicClient{sdk: c}, nil
}
func (c *AnthropicClient) Provider() string { return "anthropic" }
// toSDKMessages 把中立 Message 转 SDK 类型。
// image/document 的 Data 字段为原始字节,此处进行 base64 编码。
func toSDKMessages(msgs []Message) []anthropic.MessageParam {
out := make([]anthropic.MessageParam, 0, len(msgs))
for _, m := range msgs {
blocks := make([]anthropic.ContentBlockParamUnion, 0, len(m.Blocks))
for _, b := range m.Blocks {
switch b.Type {
case "text":
blocks = append(blocks, anthropic.NewTextBlock(b.Text))
case "image":
// NewImageBlockBase64 期望 base64 字符串,Data 字段存原始字节
encoded := base64.StdEncoding.EncodeToString(b.Data)
blocks = append(blocks, anthropic.NewImageBlockBase64(b.MimeType, encoded))
case "document":
// NewDocumentBlock 泛型支持 Base64PDFSourceParam,Data 字段存原始字节
encoded := base64.StdEncoding.EncodeToString(b.Data)
blocks = append(blocks, anthropic.NewDocumentBlock(anthropic.Base64PDFSourceParam{
Data: encoded,
}))
}
}
role := anthropic.MessageParamRoleUser
if m.Role == RoleAssistant {
role = anthropic.MessageParamRoleAssistant
}
out = append(out, anthropic.MessageParam{Role: role, Content: blocks})
}
return out
}
func (c *AnthropicClient) Stream(ctx context.Context, modelID string, req Request) (<-chan StreamEvent, error) {
params := anthropic.MessageNewParams{
Model: anthropic.Model(modelID),
MaxTokens: int64(req.MaxOutputTokens),
Messages: toSDKMessages(req.Messages),
}
if req.SystemPrompt != "" {
params.System = []anthropic.TextBlockParam{{Text: req.SystemPrompt}}
}
if req.Reasoning {
params.Thinking = anthropic.ThinkingConfigParamOfEnabled(int64(req.MaxOutputTokens / 2))
}
stream := c.sdk.Messages.NewStreaming(ctx, params)
out := make(chan StreamEvent, 16)
// send 在 ctx 取消时立即放弃发送,防止下游 reader 已退出导致 goroutine 永久阻塞。
send := func(ev StreamEvent) bool {
select {
case <-ctx.Done():
return false
case out <- ev:
return true
}
}
go func() {
defer close(out)
var inputTokens int
for stream.Next() {
event := stream.Current()
// SDK v1.38: AsAny() returns unexported anyMessageStreamEvent which cannot
// be used in a type switch. Use event.Type string + typed As*() methods.
switch event.Type {
case "message_start":
ev := event.AsMessageStart()
inputTokens = int(ev.Message.Usage.InputTokens)
case "content_block_delta":
ev := event.AsContentBlockDelta()
switch ev.Delta.Type {
case "text_delta":
if text := ev.Delta.AsTextDelta().Text; text != "" {
if !send(StreamEvent{Type: EventTextDelta, Text: text}) {
return
}
}
case "thinking_delta":
if thinking := ev.Delta.AsThinkingDelta().Thinking; thinking != "" {
if !send(StreamEvent{Type: EventThinkingDelta, Text: thinking}) {
return
}
}
}
case "message_delta":
ev := event.AsMessageDelta()
if !send(StreamEvent{Type: EventUsage, Usage: &Usage{
PromptTokens: inputTokens,
CompletionTokens: int(ev.Usage.OutputTokens),
}}) {
return
}
case "message_stop":
if !send(StreamEvent{Type: EventDone, Reason: "end_turn"}) {
return
}
}
}
if err := stream.Err(); err != nil {
send(StreamEvent{Type: EventError, Err: fmt.Errorf("anthropic stream: %w", err)})
}
}()
return out, nil
}
func (c *AnthropicClient) CountTokens(ctx context.Context, modelID string, req Request) (int, error) {
params := anthropic.MessageCountTokensParams{
Model: anthropic.Model(modelID),
Messages: toSDKMessages(req.Messages),
}
if req.SystemPrompt != "" {
// MessageCountTokensParams.System 是 MessageCountTokensParamsSystemUnion,
// 不同于 MessageNewParams.System ([]TextBlockParam)
params.System = anthropic.MessageCountTokensParamsSystemUnion{
OfTextBlockArray: []anthropic.TextBlockParam{{Text: req.SystemPrompt}},
}
}
res, err := c.sdk.Messages.CountTokens(ctx, params)
if err != nil {
return 0, fmt.Errorf("anthropic count_tokens: %w", err)
}
return int(res.InputTokens), nil
}